如何用Numpy替代Python循环实现重复项系数累加？

完全可以用Numpy的内置矢量操作实现，比Python循环字典的方式效率高几个数量级，尤其适合你这种百万级别的大数组场景。以下是两种适配不同情况的实现方案：

方案一：通用高效方案（推荐） #

该方案基于排序+分段求和，适用于唯一元素数量少、元素取值范围大的场景（正好匹配你的需求：唯一元素数量约为√N）：

import numpy as np

# 示例输入
items = np.array([1,1,2,2,3,4], dtype=np.int64)
coefficients = np.array([1,2,3,4,5,6], dtype=np.float64)

# 1. 按items排序，同时同步排序coefficients
sorted_indices = np.argsort(items)
sorted_items = items[sorted_indices]
sorted_coeffs = coefficients[sorted_indices]

# 2. 获取唯一元素及其在排序后数组中的起始位置
unique_items, unique_positions = np.unique(sorted_items, return_index=True)

# 3. 对排序后的系数按唯一元素的分段位置求和
unique_coefficients = np.add.reduceat(sorted_coeffs, unique_positions)

执行后得到的结果：

• unique_items → array([1, 2, 3, 4])
• unique_coefficients → array([3., 7., 5., 6.])

原理说明： #

• 排序让相同元素连续排列，避免了字典查找的开销；
• np.unique直接获取唯一元素和分段起始位置，是C级别的高效操作；
• np.add.reduceat对连续的系数段批量求和，比Python循环累加快得多。

方案二：小范围整数元素优化方案 #

如果你的items元素是取值范围较小的非负整数，可以用np.bincount进一步简化代码：

import numpy as np

# 仅当items为非负整数且最大值不大时适用
unique_items = np.unique(items)
unique_coefficients = np.bincount(items, weights=coefficients)[unique_items]

注意： #

如果items的最大值远大于唯一元素数量（比如元素是百万级别的整数但唯一元素只有几千），这种方法会创建一个巨大的中间数组，浪费内存，此时推荐用方案一。

性能对比 #

你的原Python循环字典方法，每个元素都要执行字典的查找、更新操作，属于Python级别的循环，对于1e7规模的数据会非常缓慢；而上述Numpy方案全部是C级别的批量操作，能将耗时降低到原方法的1%甚至更低。

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