Python Numpy如何根据第二个索引数组对目标数组选择性求和

直接使用NumPy内置的np.bincount方法即可实现，无需手动编写循环，底层为C语言实现，运算效率远高于Python层级的遍历逻辑。

实现逻辑 #

np.bincount支持传入权重参数，会自动按输入的索引值分组，对对应位置的权重值求和，刚好匹配需求：

• 把索引数组B和值数组A都展平为一维
• 以展平后的B作为索引，展平后的A作为权重传入np.bincount，直接得到分组求和结果

代码示例 #

示例1：基础场景 #

import numpy as np

# 输入数组
A = np.array([[1, 2], [6, 10]])
B = np.array([[0, 1], [1, 0]])

# 核心实现
C = np.bincount(B.flatten(), weights=A.flatten()).astype(int)
print(C)
# 输出：[11  8]，和预期完全一致

示例2：全部分到同一组的场景 #

# 更换B数组为全1
B = np.array([[1, 1], [1, 1]])

C = np.bincount(B.flatten(), weights=A.flatten()).astype(int)
print(C)
# 输出：[ 0 19]，和补充说明的预期完全一致

适配说明 #

• 该实现天然支持任意维度的A、B数组，只要两者形状完全相同，仅需展平即可，不需要修改其他逻辑
• 如果需要固定输出数组C的长度，可添加minlength参数指定，比如np.bincount(B.flatten(), weights=A.flatten(), minlength=10)会生成长度为10的结果数组，未出现的索引位置值为0
• 注意你提供的手动循环示例里C = np.zeros(np.max(B))存在边界问题：如果B的最大值为1，np.max(B)返回1，生成的C长度为1只有索引0，赋值到索引1时会报错，正确写法应为C = np.zeros(np.max(B) + 1)，np.bincount会自动处理这个边界问题。

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