python - 具有来自另一个数组的索引的 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 的子数组。例如a[0] 为 1，所以我们选择 u[0,1]，a[1] 为 0，所以我们选择 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可能会有帮助。