按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列，第三列是各种分数。

我想快速把所有的分数加起来，并按照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

我注意到有提到 numpy 这个标签，不过如果你不介意使用 pandas（或者你是通过这个模块来读取这些数据），那么这个任务就可以用一行代码搞定：

import pandas as pd

df = pd.DataFrame({'id': [1,1,1,2,2,2,3], 'score': [20,30,0,4,8,9,18]})

所以你的数据表看起来会是这样的：

  id  score
0   1     20
1   1     30
2   1      0
3   2      4
4   2      8
5   2      9
6   3     18

现在你可以使用 groupby() 和 sum() 这两个函数：

df.groupby(['id'], sort=False).sum()

这样就能得到你想要的结果：

    score
id       
1      50
2      21
3      18

默认情况下，数据表会被排序，所以我使用了 sort=False 这个选项，这样在处理非常大的数据表时可能会提高速度。

Answer 2

如果你只是想要计算总和，那你可以使用 bincount。如果你还需要其他的分组操作，比如乘积、平均值、标准差等等，可以看看这个链接：https://github.com/ml31415/numpy-groupies。这是目前最快的 Python/Numpy 分组操作，里面有速度对比的资料。

你在这里的总和操作看起来会是这样的：

res = aggregate(id, score)

Answer 3

你可以使用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。