Numpy：高效的矩阵索引

Question

我有两个 numpy arrays：

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

其中k和n是整数。

B的系数是矩阵A的索引。 例如， 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]但后者适用于具有任意维度的 arrays，只要您的索引位于最后一个维度。

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