Numpy: How to index 2d array with 1d array?

Question

I have a 2d array:

a = np.random.randint(100, size=(6, 4))
[[72 76 40 11]
 [48 82  6 87]
 [53 24 25 99]
 [ 7 94 82 90]
 [28 81 10  9]
 [94 99 67 58]]

And a 1d array:

idx = np.random.randint(4, size=6)
[0, 3, 2, 1, 0, 2]

Is it possible to index the 2d array so that the result is:

a[idx]
[72, 87, 25, 94, 28, 67]

Answer 1

Since you have the column indices, all you need are the row indices. You can generate those with arange .

>>> a[np.arange(len(a)), idx]
 array([72, 87, 25, 94, 28, 67])

Answer 2

Is there any way to get by this without arange ? It seems counterintuitive to me that something like

a[idx.reshape(-1,1)]

or

a[:,idx]

would not produce this result.

Answer 3

You can also use np.diagonal if you want to avoid np.arrange.

 a = np.array([[72, 76, 40, 11],
              [48, 82,  6, 87],
              [53, 24, 25, 99],
              [ 7, 94, 82, 90],
              [28, 81, 10,  9],
              [94, 99, 67, 58]])

idx = np.array([0, 3, 2, 1, 0, 2])

index into each array in 2d array using idx

>>> a_idx = a[...,idx]
>>> a_idx

array([[72, 11, 40, 76, 72, 40],
       [48, 87,  6, 82, 48,  6],
       [53, 99, 25, 24, 53, 25],
       [ 7, 90, 82, 94,  7, 82],
       [28,  9, 10, 81, 28, 10],
       [94, 58, 67, 99, 94, 67]])

diagonal is where the position of idx and each array in 2d array line up

>>> np.diagonal(a_idx)

array([72, 87, 25, 94, 28, 67])