如何按另一数组对应值对numpy数组元素分组求和？

问题描述 #

给定两个numpy数组：

a = np.array([2,3,4,5,6,7,8,9,10])
b = np.array([1,1,1,2,3,4,5,5,6])

需要按照b中相同的值对a的对应元素分组求和，例如b中值为1的元素对应a的前3个元素求和（2+3+4=9），值为5的元素对应a的8和9求和（8+9=17），最终得到结果：

sum1 = np.array([9,5,6,7,17,10])

原循环代码的问题 #

你编写的循环逻辑存在两处关键错误：

1. 循环范围range(1, len(a)-1)会漏掉数组末尾的元素，无法处理最后一组数据；
2. 仅对相邻的同组元素求和，没有累加整个分组的所有元素，导致多元素分组的结果错误。

高效实现方法（推荐） #

对于大数据量（数千条数据），优先使用numpy的向量化操作，避免低效的循环。以下是最优实现：

import numpy as np

a = np.array([2,3,4,5,6,7,8,9,10])
b = np.array([1,1,1,2,3,4,5,5,6])

# 获取b的唯一值，以及每个元素在唯一值数组中的索引
unique_vals, group_indices = np.unique(b, return_inverse=True)
# 按分组索引对a的元素求和
sum_result = np.bincount(group_indices, weights=a)

print(sum_result)  # 输出: [ 9.  5.  6.  7. 17. 10.]

逻辑说明 #

• np.unique(b, return_inverse=True)：返回b中的唯一值数组，以及b每个元素对应到唯一值数组的索引（用于标记分组）；
• np.bincount(group_indices, weights=a)：根据分组索引，对a中对应元素求和，自动完成分组累加，效率远高于循环。

如果需要保持结果与b中分组首次出现的顺序一致（当b的唯一值无序时），可以调整为：

# 获取分组首次出现的索引并排序
unique_first_indices = np.unique(b, return_index=True)[1]
sorted_indices = np.sort(unique_first_indices)
# 按首次出现顺序排列求和结果
sum_result_ordered = sum_result[np.argsort(unique_first_indices)]

循环实现（不推荐用于大数据） #

如果一定要用循环，修正后的逻辑如下：

sum1 = []
current_sum = a[0]
current_group = b[0]

# 从第二个元素开始遍历
for num_a, num_b in zip(a[1:], b[1:]):
    if num_b == current_group:
        current_sum += num_a
    else:
        sum1.append(current_sum)
        current_sum = num_a
        current_group = num_b
# 加入最后一组的求和结果
sum1.append(current_sum)

sum1 = np.array(sum1)
print(sum1)  # 输出: [ 9  5  6  7 17 10]

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