# Indexing arrays with other arrays

Since I can't seem to fully internalise how `numpy`

advanced indexing works, here's the solution to a common indexing task, written out for future reference.

### Task

Given an array `a`

of dimension `[n, m, l]`

and an index array `idx`

of dimension `[n, k <= m]`

, I'd like to define a new array `b`

such that `b[i, j, :] = a[i, idx[i, j], :]`

.

### Example

Let's make an example and a naive solution:

```
import numpy as np
a = np.random.random((3, 3, 2))
idx = np.array([[1, 2], [0, 2], [0, 1]])
naive = np.zeros((*idx.shape, 2))
for i in range(idx.shape[0]):
for j in range(idx.shape[1]):
naive[i, j] = a[i, idx[i, j]]
```

### Solution

The indexing-based solution, from this blog post is:

```
b = a[np.arange(a.shape[0])[:, None], idx, :]
# verify solution
np.testing.assert_array_equal(naive, b)
```

This also works with `torch`

tensors and seems reasonably fast. If I notice any performance problems, I'll update this post with a (hopefully) more efficient version...