Есть ли функция NumPy для вычисления частичных сумм?

Question

Учитывая массив 1-D, A и список n уникальных и отсортированных индексов idx, я хотел бы вычислить суммы от A[idx[i]:idx[i + 1]] для i = 0 до n - 1. Решение на основе:

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.