python - 使用2D矩阵作为numpy中3D矩阵的索引吗？

假设我有一个形状为2x3x3的数组，这是一个3D矩阵。我也有一个形状为3x3的2D矩阵，我希望将其用作沿第一轴的3D矩阵的索引。示例如下。

示例运行:

>>> np.random.randint(0,2,(3,3)) # index
array([[0, 1, 0],
       [1, 0, 1],
       [1, 0, 0]])

>> np.random.randint(0,9,(2,3,3)) # 3D matrix
array([[[4, 4, 5],
        [2, 6, 7],
        [2, 6, 2]],

       [[4, 0, 0],
        [2, 7, 4],
        [4, 4, 0]]])
>>> np.array([[4,0,5],[2,6,4],[4,6,2]]) # result
array([[4, 0, 5],
       [2, 6, 4],
       [4, 6, 2]])

最佳答案

看来您正在使用2D数组作为索引数组和3D数组来选择值。因此，您可以使用NumPy的 advanced-indexing -

# a : 2D array of indices, b : 3D array from where values are to be picked up
m,n = a.shape
I,J = np.ogrid[:m,:n]
out = b[a, I, J] # or b[a, np.arange(m)[:,None],np.arange(n)]

如果您打算使用a索引到最后一个轴，只需将a移动到那里:b[I, J, a]。

sample 运行-

>>> np.random.seed(1234)
>>> a = np.random.randint(0,2,(3,3))
>>> b = np.random.randint(11,99,(2,3,3))
>>> a  # Index array
array([[1, 1, 0],
       [1, 0, 0],
       [0, 1, 1]])
>>> b  # values array
array([[[60, 34, 37],
        [41, 54, 41],
        [37, 69, 80]],

       [[91, 84, 58],
        [61, 87, 48],
        [45, 49, 78]]])
>>> m,n = a.shape
>>> I,J = np.ogrid[:m,:n]
>>> out = b[a, I, J]
>>> out
array([[91, 84, 37],
       [61, 54, 41],
       [37, 49, 78]])