Numpy:高效的矩阵索引

Question

我有两个numpy数组：

• A的形状为(k, 4, 2)
• B的形状为(n, 12, 2)

其中k和n是整数。

CCD_ 7的系数是矩阵CCD_。例如，B[0, 1, :]是一个数组[k, l]，并且我们确信A[k , l]存在。

我想做的是构建一个与B大小相同的矩阵C，这样，对于所有i, j，C[i, j] = A[B[i, j, 0], B[i, j, 1]]

有有效的方法吗？

我尝试过A[:, B]、A[0, B]之类的东西，但都不成功。我也可以用for循环来实现，但我认为用numpy实现会更快。

对于那些想尝试的人，我准备了一个小的启动代码(用k=n=2测试方法：

import numpy as np
a = np.array([[
[73,  -25],
[97, -25],
[73,  107],
[97, 107]],
[[81,  43],
[86, 43],
[81,  50],
[86, 43]]
])
b = np.array(
[[[0, 2],
[0, 0],
[0, 3],
[0, 1],
[1, 0],
[1, 2],
[1, 1],
[1, 3],
[0, 0],
[0, 0],
[0, 0],
[0, 0]],
[[0, 3],
[0, 2],
[0, 1],
[0, 0],
[1, 0],
[1, 2],
[1, 1],
[1, 3],
[0, 0],
[0, 0],
[0, 0],
[0, 0]]]
)
#the answer should be :
c = np.array(
[[[73,  107],
[73,  -25],
[97, 107],
[97, -25],
[81,  43],
[81,  50],
[86, 43],
[86, 50],
[73,  -25],
[73,  -25],
[73,  -25],
[73,  -25]],
[[97, 107],
[73,  107],
[97, -25],
[73,  -25],
[81,  43],
[81,  50],
[86, 43],
[86, 50],
[73,  -25],
[73,  -25],
[73,  -25],
[73,  -25]],
]
)

希望你清楚，提前感谢

Answer 1

a[b[...,0],b[...,1]]

应该做预期的事情。您可以使用b[:, :, 0]而不是b[..., 0]，但后者适用于具有任意维度的数组，只要您的索引位于最后一个维度即可。

Answer 2

你几乎做到了：

C=A[B[:,:,0],B[:,:,1]]

但你让我困惑了一分钟，因为你的测试c有几个错误的值：

>>> a,b,c = ...
>>> C = a[b[:,:,0],b[:,:,1]]
>>> np.all([C[i,j]==a[b[i,j,0],b[i,j,1]] for i in range(2) for j in range(12)])
True
>>> C==c
array([[[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True, False],
[ True, False],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True]],
[[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True],
[ True, False],
[ True, False],
[ True,  True],
[ True,  True],
[ True,  True],
[ True,  True]]])