最有效的方法来汇总巨大的2D NumPy数组,按ID列分组？

Question

我有一个海量数据阵列(500k行),看起来像:

id  value  score
1   20     20
1   10     30
1   15     0
2   12     4
2   3      8
2   56     9
3   6      18
...

如您所见,左侧有一个非唯一ID列,第三列中有各种分数.

我希望快速添加按ID分组的所有分数.在SQL中,这看起来像SELECT sum(score) FROM table GROUP BY id

使用NumPy,我尝试迭代每个ID,按每个ID截断表,然后将该表的得分相加.

table_trunc = table[(table == id).any(1)]
score       = sum(table_trunc[:,2])

不幸的是,我发现第一个命令是狗慢.有没有更有效的方法来做到这一点？

Answer 1

你可以使用bincount():

import numpy as np

ids = [1,1,1,2,2,2,3]
data = [20,30,0,4,8,9,18]

print np.bincount(ids, weights=data)

输出为[0. 50. 21. 18.],这意味着id == 0的总和为0,id == 1的总和为50.