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

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:

1. 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.
2. 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)
```

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]])
```