numpy数组分箱求和函数咨询：大数组线性时间解决方案

当然有啦！numpy里正好有高效的方案能满足你的分箱求和需求，而且完全是线性时间复杂度，特别适配你提到的大规模、多数量数组的场景。

先明确你的核心需求：将长度为n的数组（n能被分箱数b整除）拆成b个等长的连续子数组，每个子数组求和得到结果数组。下面给你两种实用的方法：

方法一：Reshape + Sum（简洁首选） #

这种方法利用numpy的reshape（重塑数组形状）和sum（求和）操作，代码简洁且性能拉满：

示例1：数组[1,0,5,1]分2箱 #

import numpy as np
L = np.array([1, 0, 5, 1])
b = 2
step = len(L) // b  # 每个分箱的长度，这里是2
Lbin = L.reshape(b, step).sum(axis=1)
print(Lbin)  # 输出：array([1, 6])

示例2：数组[1,3,5,2,6,7]分3箱 #

L = np.array([1, 3, 5, 2, 6, 7])
b = 3
step = len(L) // b  # 每个分箱的长度，这里是2
Lbin = L.reshape(b, step).sum(axis=1)
print(Lbin)  # 输出：array([ 4,  7, 13])

性能说明：

• reshape在数组是连续内存布局时，只是修改了数组的视图，不会额外占用内存，时间复杂度为O(1)；
• sum操作是线性遍历数组求和，时间复杂度为O(n)；
• 整体是线性时间，完全符合你的性能要求。

方法二：numpy.add.reduceat（灵活适配非连续数组） #

如果你的数组是经过切片、转置等操作后的非连续数组，reshape可能会触发内存复制，这时候可以用numpy.add.reduceat来实现，同样是线性时间：

L = np.array([1, 3, 5, 2, 6, 7])
b = 3
step = len(L) // b
split_points = np.arange(0, len(L), step)  # 分箱的起始索引：[0,2,4]
Lbin = np.add.reduceat(L, split_points)
print(Lbin)  # 输出：array([ 4,  7, 13])

这个方法通过指定分箱的起始索引，直接对数组进行分段求和，无需修改数组形状，灵活性更强。

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