按ID列分组，求解巨大的二维NumPy数组的最有效方法是什么？

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列，第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])

很不幸，我发现第一个命令非常慢。有没有更有效的方法来完成这个任务？

- thegreatt

---

请参考 https://dev59.com/w2445IYBdhLWcg3w4-Be 了解有关numpy分组的信息。 - agf

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。

- HYRY

---

但这仅适用于1维数组。海报想要一个2维数组的解决方案。有没有解决方案？ - pfc

Answer 2

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

- piRSquared

Answer 3

我注意到了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，这可能有助于处理大型数据框的速度。

- Cleb

Answer 4

如果你只需要sum，那么你可能想使用bincount。如果你还需要其他分组操作，比如乘积、平均值、标准差等，请查看https://github.com/ml31415/numpy-groupies。这是目前最快的Python / NumPy分组操作，可以在那里查看速度比较。

你的求和操作应该像这样：

res = aggregate(id, score)

- Michael

Answer 5

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

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与数据表的总体大小相匹配，那么显然会遇到麻烦。

- dtlussier

Answer 6

numpy_indexed 包具有矢量化功能，可以高效地执行此操作，除此之外还有许多相关的操作：

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

- Eelco Hoogendoorn

Answer 7

也许可以使用 itertools.groupby，按照 ID 进行分组，然后迭代分组数据。（数据必须按照分组函数排序，即按照 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)

- roqvist

---

我认为这不太可能很快，因为它是在Python中完成工作，而不是像在numpy中使用C一样。 - agf