如何用Numpy的add.reduceat实现分组求和以优化伯努利试验计算

问题背景 #

需要优化一段Numpy代码，功能是对数组num中的每个元素n（代表对应桶的物品数量），执行n次伯努利试验（成功概率为qEff），统计每个桶的试验成功次数。原代码通过循环实现效率较低，计划先生成所有试验结果数组rol，再按num的元素分组求和，但不知道如何构建np.add.reduceat所需的indices参数。

原循环代码：

import numpy as np
rng = np.random.default_rng()
num = np.fromiter((rng.choice([0, 1], size=n, p=[1-qEff, qEff]).sum()
                   for n in num), dtype='int')

优化思路框架代码：

rol = rng.choice([0, 1], size=num.sum(), p=[1-qEff, qEff])
num = some_smart_sum(rol, num)

解决方案 #

要使用np.add.reduceat实现高效分组求和，核心是生成分组起始索引数组，具体步骤如下：

1. 计算前缀和：用np.cumsum(num)得到每个桶的累计元素结束位置。例如num=[2,3,4]，前缀和为[2,5,9]，表示前2个元素属于第一个桶，前5个属于前两个桶，以此类推。
2. 构建起始索引：在前缀和的前n-1个元素前添加0，得到每个分组的起始位置。对于上面的例子，起始索引为[0,2,5]，对应rol[0:2]、rol[2:5]、rol[5:9]三个分组。

完整优化代码

import numpy as np
rng = np.random.default_rng()

# 示例参数，根据实际场景替换
qEff = 0.3
num = np.array([2, 3, 4, 1])

# 生成所有试验结果
rol = rng.choice([0, 1], size=num.sum(), p=[1 - qEff, qEff])

# 构建reduceat所需的indices参数
indices = np.r_[0, np.cumsum(num)[:-1]]

# 分组求和得到每个桶的成功次数
num = np.add.reduceat(rol, indices)

关键代码解释

• np.cumsum(num)[:-1]：取前缀和的前n-1个元素，作为后续分组的起始位置（跳过最后一个累计值，因为reduceat会自动处理到数组末尾）。
• np.r_[0, ...]：快速拼接起始索引0，形成完整的分组起始数组。
• np.add.reduceat(rol, indices)：按indices指定的位置对rol进行分段求和，效率远高于循环实现。

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