Numpy：高效的矩阵索引

Question

我有两个 numpy 数组：

• 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]，但后者适用于任意维度的数组，只要您的索引位于最后一维。

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