Numpy index nested array with indices from another array

Question

I have a numpy array like:

u = np.arange(10).reshape(5, 2)

array([[0, 1],
   [2, 3],
   [4, 5],
   [6, 7],
   [8, 9]])

I have a second array like

a = np.array([1,0,0,1,0])

I would like to use the values from a to index the subarrays of u. E.g. a[0] is 1, so we chose u[0,1], a[1] is 0, so we choose u[1, 0] and so forth.

I have tried lots of things, and would like to do it without for loops. Even after reading numpys indexing guide I have not really found how to do it.

Things that I have tried that failed:

>>> u[:, [0,0,1,0,1]]
array([[0, 0, 1, 0, 1],
   [2, 2, 3, 2, 3],
   [4, 4, 5, 4, 5],
   [6, 6, 7, 6, 7],
   [8, 8, 9, 8, 9]])

u[[True, False, True, True, True]]
array([[0, 1],
   [4, 5],
   [6, 7],
   [8, 9]])

Lastly to clear up confusions, here is what I want, however with python loops:

>>> x = []
>>> ct = 0
>>> for i in u:
        x.append(i[a[ct]])
        ct += 1

>>> x
[1, 2, 4, 7, 8]

Thanks in advance.

Answer 1

Use:

import numpy as np

u = np.arange(10).reshape(5, 2)
a = np.array([1, 0, 0, 1, 0])
r, _ = u.shape  # get how many rows to use in np.arange(r)

print(u[np.arange(r), a])

Output

[1 2 4 7 8]

For more on indexing, you can read the documentation and also this article could be helpful.