如何避免循环，高效实现Numpy数组按维度索引求和？

假设有两个NumPy数组A和B，A中的值对应B第一维度的索引：

import numpy as np

A = np.array([[ 0,  3,  4], 
              [ 5,  6,  0],
              [ 0,  1,  9],
              [11, 12, 13]])

B = np.array([[0, 1, 1],
              [1, 2, 0],
              [0, 2, 3],
              [3, 3, 3],
              [0, 3, 1],
              [3, 1, 1],
              [1, 1, 0]])

需求为：对B的每一行j，计算所有A[B[j, i], i]的和，得到形状为7的数组。目前已通过嵌套循环实现：

output_array = np.zeros(B.shape[0])
for i in range(A.shape[1]):
    for j in range(B.shape[0]):
        output_array[j] += A[B[j, i], i]

运行后得到结果：

array([ 6., 10., 14., 36., 12., 17., 15.])

针对大型数组，是否存在无需显式循环的快速高效实现方式？曾找到类似方案，但无法扩展到该求和场景。

高效矢量化实现方案 #

可以利用NumPy的高级索引直接提取所有目标元素，再按行求和，完全规避Python层循环，性能远优于嵌套循环：

import numpy as np

# 生成与B列数匹配的列索引数组，通过广播适配B的形状
col_indices = np.arange(B.shape[1])[np.newaxis, :]

# 一次性提取所有A[B[j,i], i]元素，得到(7,3)形状的数组
selected_elements = A[B, col_indices]

# 按行求和得到目标数组
output_array = selected_elements.sum(axis=1)

验证结果：

print(output_array)
# 输出：array([ 6, 10, 14, 36, 12, 17, 15])

原理说明 #

• col_indices生成了[0,1,2]的行向量，通过广播自动扩展为与B相同的(7,3)形状，确保每个B[j,i]都能对应到正确的列索引i；
• A[B, col_indices]利用NumPy的高级索引机制，一次性定位所有需要的元素，避免了逐元素循环；
• 最终调用sum(axis=1)对每行元素求和，直接得到目标结果。

这种实现完全依赖NumPy底层的C级矢量化操作，在处理大型数组时，性能会比嵌套循环提升几个数量级。

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