python - 对按 ID 列分组的巨大 2D NumPy 数组求和的最有效方法？

Question

我有一个海量数据数组(50 万行)，如下所示:

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 列，第 3 列中有各种分数。

我希望快速将所有分数相加，并按 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。