如何在python中求和数组的特定元素

Question

有一个简单的操作，我需要对非常大的数组(数千个元素(做几百次，所以我需要最有效的解决方案(对于循环太慢(：

我有两个阵列，例如

a = np.array([1,23,25,100,100,101])
b = np.array([1,2,2,3,4,4])

我现在想得到a中所有元素的和，其中b的值相同。即

[1,48,100,201]

我可以做：

#first index of each unique entry in b
u = np.unique(b,return_index=True)[1]
#split array and sum
list(map(sum, np.split(a,u[1:])))

但这有点慢，而且只有对b中的条目进行排序时才有效。有其他方法吗？

Answer 1

尝试：

>>> [a[b==n].sum() for n in np.unique(b)]
[1, 48, 100, 201]

如果您对使用pandas:持开放态度

>>> pd.DataFrame({"a": a, "b": b}).groupby("b").sum()["a"].tolist()
[1, 48, 100, 201]