python - 是否有计算部分和的 NumPy 函数？

Question

给定一个一维数组 A 和一个 n 唯一和排序索引 idx 的列表，我想计算的总和A[idx[i]:idx[i + 1]] 从 i = 0 到 n - 1。基于 for 的解决方案是:

S = [A[idx[i]:idx[i + 1]].sum() for i in range(n - 1)]

但我想如果 n 很大，因为它是在 Python 级别完成的，那么这会很慢。是否有 NumPy 函数可以实现此目的(希望更快)？

Answer 1

您正在寻找 reduceat np.add的方法:

np.add.reduceat(A, idx)[:-1]

[:-1] 只是为了删除 reduceat 附加的最后一个元素以及来自 idx[-1] 的总和到 A.size。