有没有一个NumPy函数可以计算部分和？

Question

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

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

但是，如果n很大，这样做在Python层面上可能会非常缓慢。是否存在NumPy函数可以实现这一点（希望更快）？

- MikeL

---

谢谢您的提醒。我刚刚清理了一些旧问题。 - MikeL

Answer 1

你正在寻找 reduceat 方法的相关信息，该方法与 np.add 相关：

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

< p > [:-1] 的作用是移除 reduceat 附加到 idx[-1] 到 A.size 的总和的最后一个元素。

- Mad Physicist

---

谢谢！我之前不知道这个方法。 - MikeL

然而，我认为应该小心处理 idx[-1] == len(A) 这种情况（在 reducecat 中会引发 IndexError）。因此，也许需要单独处理这种特殊情况。 - MikeL