Numpy: Как проиндексировать 2d массив с 1d массивом?

Question

У меня есть 2d массив:

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

И 1d массив:

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

Можно ли индексировать 2d массив так, чтобы результат был:

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

Answer 1

Поскольку у вас есть индексы столбцов, все, что вам нужно, это индексы строк. Вы можете генерировать их с помощью arange.

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