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

时间:2011-08-17 07:45:28

标签: python 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中,这看起来像SELECT sum(score) FROM table GROUP BY id

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

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

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

7 个答案:

答案 0 :(得分:10)

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

答案 1 :(得分:1)

我注意到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

默认情况下,数据框将被排序,因此我使用标记{{1}},这可能会提高大型数据帧的速度。

答案 2 :(得分: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.any更有效,但如果您拥有大量唯一ID以及数据表的大整体大小,则显然会遇到麻烦。

答案 3 :(得分:0)

如果您只查找sum,则可能需要使用bincount。如果您还需要其他分组操作,例如product,mean,std等,请查看https://github.com/ml31415/numpy-groupies。它是最快的python / numpy分组操作,请参阅那里的速度比较。

你的总和操作如下:

res = aggregate(id, score)

答案 4 :(得分:0)

numpy_indexed包具有矢量化功能,可以有效地执行此操作,此外还有许多相关的操作:

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

答案 5 :(得分:0)

您可以使用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]

然后:

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

答案 6 :(得分:-1)

也许使用itertools.groupby,您可以对ID进行分组,然后迭代分组数据。

(数据必须按照func分组,在本例中为ID)

>>> data = [(1, 20, 20), (1, 10, 30), (1, 15, 0), (2, 12, 4), (2, 3, 0)]
>>> groups = itertools.groupby(data, lambda x: x[0])
>>> for i in groups:
        for y in i:
            if isinstance(y, int):
                print(y)
            else:
                for p in y:
                    print('-', p)

输出:

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