Python 对按ID列分组的巨大2D NumPy数组求和的最有效方法是什么？

我有一个海量数据阵列（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中，这看起来像是

selectsum（score）fromtablegroupbyid

使用NumPy，我尝试遍历每个ID，按每个ID截断表，然后将该表的分数相加

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

不幸的是，我发现第一个命令是dog slow。有更有效的方法吗？

您可以使用bincount（）

---

输出是[0.50.21.18.]，这意味着id==0的和是0，id==1的和是50。

可能使用

itertools.groupby

，您可以在id上分组，然后在分组的数据上迭代

（数据必须按func分组排序，在本例中为ID）

输出：

1
- (1, 20, 20)
- (1, 10, 30)
- (1, 15, 0)
2
- (2, 12, 4)
- (2, 3, 0)

---

您可以尝试使用布尔运算：

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

[((ids == i)*data).sum() for i in np.unique(ids)]

---

这可能比使用

np更有效。任何

，但如果数据表的总体大小很大，并且需要使用大量的唯一ID，则显然会遇到问题。

如果只查找

总和

，则可能需要使用

二进制计数

。如果您还需要其他分组操作，如产品、平均值、标准等，请查看。这是最快的python/numpy分组操作，请参见此处的速度比较

那里的求和运算如下所示：

res = aggregate(id, score)

---

我注意到了

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（）

：

这将为您提供所需的输出：

    score
id       
1      50
2      21
3      18

默认情况下，数据帧将被排序，因此我使用标志

sort=False

，这可能会提高大型数据帧的速度。

除了许多此类相关操作外，该软件包还具有向量化功能以高效执行此操作：

import numpy_indexed as npi
npi.group_by(id).sum(score)

---

您可以使用for循环和

numba

from numba import njit

@njit
def wbcnt(b, w, k):
    bins = np.arange(k)
    bins = bins * 0
    for i in range(len(b)):
        bins[b[i]] += w[i]
    return bins

使用@HYRY的变量

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

然后：

---

定时

%timeit wbcnt(ids, data, 4)
%timeit np.bincount(ids, weights=data)

1000000 loops, best of 3: 1.99 µs per loop
100000 loops, best of 3: 2.57 µs per loop

---

有关numpy分组的信息，请参阅。我认为这不太可能很快，因为它使用Python而不是C来完成工作，就像在

numpy

中完成一样。但这只适用于1d数组。海报需要2d阵列的解决方案。有什么解决办法吗？

from numba import njit

@njit
def wbcnt(b, w, k):
    bins = np.arange(k)
    bins = bins * 0
    for i in range(len(b)):
        bins[b[i]] += w[i]
    return bins

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

wbcnt(ids, data, 4)

array([ 0, 50, 21, 18])

%timeit wbcnt(ids, data, 4)
%timeit np.bincount(ids, weights=data)

1000000 loops, best of 3: 1.99 µs per loop
100000 loops, best of 3: 2.57 µs per loop