Slice a 3d numpy array using a 1d lookup between indices

import numpy as np a = np.arange(12).reshape(2, 3, 2) b = np.array([2, 0])

b maps i to j where i and j are the first 2 indexes of a, so a[i,j,k]

Desired result after applying b to a is:

[[4 5] [6 7]]

Naive solution:

c = np.empty(shape=(2, 2), dtype=int) for i in range(2): j = b[i] c[i, :] = a[i, j, :]

Question: Is there a way to do this using a numpy or scipy routine or routines or fancy indexing?

Application: Reinforcement Learning finite MDPs where b is a deterministic policy vector pi(a|s), a is the state transition probabilities p(s’|s,a) and c is the state transition matrix for that policy vector p(s’|s). The arrays will be large and this operation will be repeated a large number of times so needs to be scaleable and fast.

What I have tried:

- Compiling using numba but line profiler suggests my code is slower compared to a similarly sized numpy routine. Also numpy is more widely understood and used.
- Maintaining pi(a|s) as a sparse matrix (all zero except one 1 per row) b_as_a_matrix and then using einsum but this involves storing and updating the matrix and creates more work (an extra loop over j and sum operation).

c = np.einsum('ij,ijk->ik', b_as_a_matrix, a)

## Answer

Numpy arrays can be indexed using other arrays as indices. See also: NumPy selecting specific column index per row by using a list of indexes.

With that in mind, we can vectorize your loop to simply use `b`

for indexing:

>>> import numpy as np >>> a = np.arange(12).reshape(2, 3, 2) >>> b = np.array([2, 0]) >>> i = np.arange(len(b)) >>> i array([0, 1]) >>> a[i, b, :] array([[4, 5], [6, 7]])