使用索引数组访问 numpy 数组

Question

我正在尝试使用另一个在每个位置提供索引的数组 B 来访问一个 numpy 数组 A：

A = np.array([[1,2],[3,4]])
B = np.array([[[0,0],[0,0]],[[0,1],[0,1]]])

期望的输出：

C = array([[1,1],[3,3]])

我还没有使用 np.take() 或高级索引让它工作。我可以迭代地完成它，但我的数组大约是 10**7 所以我希望有一种更快的方法。

Answer 1

我可能应该坚持先看迭代解决方案，但这是数组一：

In [45]: A[B[:,:,1], B[:,:,0]]
Out[45]: 
array([[1, 1],
       [3, 3]])

我先试了A[B[:,:,0], B[:,:,1]]，内维度的自然顺序。你自己的代码本可以让我省去那次审判。

advanced indexing 的关键是您必须为每个维度创建或定义单独的数组（可广播）。我们可以将该索引视为一个元组：

 idx = (B[:,:,0], B[:,:,1])
 A[idx]

Answer 2

在@hpaulj 上添加另一种方法是：

idx =  tuple(B[:,:,[1,0]].transpose(2,0,1))
A[idx]
# array([[1, 1], [3, 3]])

Accessing a numpy array with an array of indices