NumPy如何通过两个2D索引矩阵索引二维矩阵得到三维结果

需求说明 #

现有6×6矩阵A，两个2×2索引矩阵B、C，要求按行配对B和C的元素：B的每一行作为A的行索引组，C的每一行作为对应A的列索引组，最终输出维度为2×2×2的三维结果矩阵。

示例输入 #

import numpy as np
A = np.arange(0,36).reshape(6,6)
B = np.array([[0,1],
              [2,4]])
C = np.array([[1,2],
              [3,4]])

期望输出 #

[[[ 1  2]
  [ 7  8]]

 [[15 16]
  [27 28]]]

实现方案 #

方案1：向量化广播索引（性能最优，无循环） #

利用Numpy的广播机制调整索引矩阵的维度后直接索引，一行代码即可完成：

result = A[B[..., np.newaxis], C[:, np.newaxis, :]]

维度调整逻辑：

• 给B新增最后一个维度，形状从(2,2)变为(2,2,1)
• 给C新增中间维度，形状从(2,2)变为(2,1,2)
• 两个索引矩阵广播后形状均为(2,2,2)，自动配对所有行、列索引组合，直接得到目标三维矩阵

方案2：列表推导式循环（易读易维护） #

如果对广播逻辑不熟悉，可以逐行遍历索引矩阵，提取子矩阵后拼接：

result = np.array([A[b_row[:, None], c_row] for b_row, c_row in zip(B, C)])

两种方案输出结果完全一致，均符合要求。 #

内容的提问来源于stack exchange，提问作者qlwang