通过检查第二个数组中的条件，在一个数组中求和元素的快速方法

Question

我有三个数组

a = array([4, 4, 3, 5, 5, 4, 6, 5, 6])
b = array([0, 1, 2, 2, 2, 3, 3, 4, 4])
c = array([0,0,0,0,0])

我想对a中b中所有元素的所有元素求和。例如，我想要跟随数组。

c = array([4,4,3+5+5,4+6,5+6])

我可以通过运行for loop执行此操作，如下所示

loop_array = scipy.unique(b)
for i in loop_array:
    c[i]=sp.sum(a[b==i])

由于我的原始a和b数组非常大，大约为百万，我无法使用for loop。有没有更快的方法来解决这个问题。如果可能，任何阵列操作都将是更好的。

Answer 1

这应该更快，但它需要更多内存用于辅助阵列

slices1 = np.insert(np.nonzero(np.diff(b))[0]+1, 0, 0)
slices2 = np.insert(np.nonzero(np.diff(b))[0]+1, len(slices1)-1, len(b))
c = [np.sum(a[j[0]:j[1]]) for j in zip(slices1, slices2)]