Numpy索引嵌套数组，索引来自另一个数组

Question

我有一个numpy数组，如：

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

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

我有第二个数组，比如

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

我想使用a中的值来索引u的子数组。 E.g. a[0] is 1, so we chose u[0,1], a[1] is 0, so we choose u[1, 0]等等

我已经尝试了很多东西，并且希望不使用for循环。即使在读了numpys indexing guide之后，我也没有真正发现如何去做

我尝试过但失败的事情：

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

最后，为了消除混淆，我希望python循环：

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

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

提前谢谢

Answer 1

使用：

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

输出

[1 2 4 7 8]

有关索引的更多信息，您可以阅读documentation，这个article可能会有所帮助