Is there any numpy group by function?

有没有一个在 numpy 中的函数将这个数组按照第一列分组？

我在网上找不到任何好的答案。

>>> a
array([[  1, 275],
[  1, 441],
[  1, 494],
[  1, 593],
[  2, 679],
[  2, 533],
[  2, 686],
[  3, 559],
[  3, 219],
[  3, 455],
[  4, 605],
[  4, 468],
[  4, 692],
[  4, 613]])

需要的产出:

array([[[275, 441, 494, 593]],
[[679, 533, 686]],
[[559, 219, 455]],
[[605, 468, 692, 613]]], dtype=object)
118587 次浏览
小开 
n = np.unique(a[:,0])
np.array( [ list(a[a[:,0]==i,1]) for i in n] )

outputs:

array([[275, 441, 494, 593], [679, 533, 686], [559, 219, 455],
[605, 468, 692, 613]], dtype=object)
小开

Numpy _ indexed软件包(免责声明: 我是它的作者)旨在填补这个缺口麻木。在数字索引中的所有操作都是完全向量化的，并且在这个库的制作过程中没有 O (n ^ 2)算法受到损害。

import numpy_indexed as npi
npi.group_by(a[:, 0]).split(a[:, 1])

注意，直接计算这些组(即 group _ by (key))上的相关属性通常更有效。(值) ，而不是先拆分成一个列表/参差不齐的数组。

小开

灵感来自 埃尔科 · 胡根多伦的图书馆，但没有他的库，并使用的事实，您的数组的第一列总是增加(如果没有，排序第一与 a = a[a[:, 0].argsort()])

>>> np.split(a[:,1], np.unique(a[:, 0], return_index=True)[1][1:])
[array([275, 441, 494, 593]),
array([679, 533, 686]),
array([559, 219, 455]),
array([605, 468, 692, 613])]

我没有“计时”((编辑)见下文) ，但这可能是更快地实现这个问题的方法:

  • 没有 Python 本地循环
  • 结果列表是 numpy 数组，如果需要对它们执行其他 numpy 操作，则不需要进行新的转换
  • 复杂性看起来是 O (n)(排序的话是 O (n log (n)))

[2021年9月编辑]我在我的 Macbook M1上运行了一个包含10k 随机整数的表。持续时间是1000通电话。

>>> a = np.random.randint(5, size=(10000, 2))  # 5 different "groups"


# Only the sort
>>> a = a[a[:, 0].argsort()]
⏱ 116.9 ms


# Group by on the already sorted table
>>> np.split(a[:, 1], np.unique(a[:, 0], return_index=True)[1][1:])
⏱ 35.5 ms


# Total sort + groupby
>>> a = a[a[:, 0].argsort()]
>>> np.split(a[:, 1], np.unique(a[:, 0], return_index=True)[1][1:])
⏱ 153.0 ms 👑


# With numpy-indexed package (cf Eelco answer)
>>> npi.group_by(a[:, 0]).split(a[:, 1])
⏱ 353.3 ms


# With pandas (cf Piotr answer)
>>> df = pd.DataFrame(a, columns=["key", "val"]) # no timer for this line
>>> df.groupby("key").val.apply(pd.Series.tolist)
⏱ 362.3 ms


# With defaultdict, the python native way (cf Piotr answer)
>>> d = defaultdict(list)
for key, val in a:
d[key].append(val)
⏱ 3543.2 ms


# With numpy_groupies (cf Michael answer)
>>> aggregate(a[:,0], a[:,1], "array", fill_value=[])
⏱ 376.4 ms

第二个场景，使用500个不同的组而不是5个。 我对熊猫感到惊讶，我跑了好几次，但它在这种情况下表现得很糟糕。

>>> a = np.random.randint(500, size=(10000, 2))


just the sort  141.1 ms
already_sorted 392.0 ms
sort+groupby   542.4 ms
pandas        2695.8 ms
numpy-indexed  800.6 ms
defaultdict   3707.3 ms
numpy_groupies 836.7 ms

我改进了答案 Ns63sr 的回答 和 Behzad Shayegh(cf 评论) 还要感谢 TMBailey注意到 argsort 的复杂性是 n log (n)。

小开

我使用了 np.unique () ，后面跟着 np.tract ()

unique = np.unique(a[:, 0:1])
answer = []
for element in unique:
present = a[:,0]==element
answer.append(np.extract(present,a[:,-1]))
print (answer)

[array([275, 441, 494, 593]), array([679, 533, 686]), array([559, 219, 455]), array([605, 468, 692, 613])]

小开

Numpy 在这里不太方便，因为所需的输出不是一个整数数组(它是一个列表对象数组)。

我建议要么用纯蟒蛇的方法。

from collections import defaultdict


%%timeit
d = defaultdict(list)
for key, val in a:
d[key].append(val)
10.7 µs ± 156 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)


# result:
defaultdict(list,
{1: [275, 441, 494, 593],
2: [679, 533, 686],
3: [559, 219, 455],
4: [605, 468, 692, 613]})

...or the pandas way:

import pandas as pd


%%timeit
df = pd.DataFrame(a, columns=["key", "val"])
df.groupby("key").val.apply(pd.Series.tolist)
979 µs ± 3.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)


# result:
key
1    [275, 441, 494, 593]
2         [679, 533, 686]
3         [559, 219, 455]
4    [605, 468, 692, 613]
Name: val, dtype: object
小开

简化 文森特 · J 的回答，并考虑 HS 星云的评论，可以使用 return_index = True而不是 return_counts = True，去掉 cumsum:

np.split(a[:,1], np.unique(a[:,0], return_index = True)[1])[1:]

输出

[array([275, 441, 494, 593]),
array([679, 533, 686]),
array([559, 219, 455]),
array([605, 468, 692, 613])]
小开

给定 X 作为要分组的项目数组，y (1D 数组)作为相应的组，下面的函数用 麻木不仁进行分组:

def groupby(X, y):
y = np.asarray(y)
X = np.asarray(X)
y_uniques = np.unique(y)
return [X[y==yi] for yi in y_uniques]

groupby(a[:,1], a[:,0])返回 [array([275, 441, 494, 593]), array([679, 533, 686]), array([559, 219, 455]), array([605, 468, 692, 613])]

小开

我们可能还会发现生成一个 dict是有用的:

def groupby(X):
X = np.asarray(X)
x_uniques = np.unique(X)
return {xi:X[X==xi] for xi in x_uniques}

Let's try it out:

X=[1,1,2,2,3,3,3,3,4,5,6,7,7,8,9,9,1,1,1]
groupby(X)
Out[9]:
{1: array([1, 1, 1, 1, 1]),
2: array([2, 2]),
3: array([3, 3, 3, 3]),
4: array([4]),
5: array([5]),
6: array([6]),
7: array([7, 7]),
8: array([8]),
9: array([9, 9])}

请注意，这本身并不是超级引人注目的-但如果我们使 X一个 object或 namedtuple，然后提供一个 groupby的功能，它会变得更有趣。稍后会放进去的。

小开

派对迟到了，但不管怎样。如果您不仅计划对数组进行分组，而且还想对它们进行求和、均值等操作，并且考虑到这样做的速度，那么您可能还需要考虑 Numpy _ groupies。所有这些组操作都进行了优化，并与 numba 进行了抖动。它们的性能很容易超过上面提到的其他解决方案。

from numpy_groupies.aggregate_numpy import aggregate
aggregate(a[:,0], a[:,1], "array", fill_value=[])
>>> array([array([], dtype=int64), array([275, 441, 494, 593]),
array([679, 533, 686]), array([559, 219, 455]),
array([605, 468, 692, 613])], dtype=object)
aggregate(a[:,0], a[:,1], "sum")
>>> array([   0, 1803, 1898, 1233, 2378])
小开

It becomes pretty apparent that a = a[a[:, 0].argsort()] is a bottleneck of all the competetive grouping algorithms, big thanks to 文森特 · J for clarifying this. Over 80% of processing time are just blown up on this argsort method and there's no easy way to replace or optimise it. numba package allows to speed up a lot of algorithms and, hopefully, argsort will attract any efforts in the future. The 剩下的部分 of grouping can be improved significantly assuming indices of first column are small.

TL;DR

大多数分组方法的 剩下的部分都包含 np.unique方法，在分组值较小的情况下，np.unique方法运算速度较慢且过量。用 np.bincount代替它效率更高，后者可以在 numba中得到改进。 关于如何改进 剩下的部分，已经取得了一些成果:

def _custom_return(unique_id, a, split_idx, return_groups):
'''Choose if you want to also return unique ids'''
if return_groups:
return unique_id, np.split(a[:,1], split_idx)
else:
return np.split(a[:,1], split_idx)


def numpy_groupby_index(a, return_groups=False):
'''Code refactor of method of Vincent J'''
u, idx = np.unique(a[:,0], return_index=True)
return _custom_return(u, a, idx[1:], return_groups)


def numpy_groupby_counts(a, return_groups=False):
'''Use cumsum of counts instead of index'''
u, counts = np.unique(a[:,0], return_counts=True)
idx = np.cumsum(counts)
return _custom_return(u, a, idx[:-1], return_groups)


def numpy_groupby_diff(a, return_groups=False):
'''No use of any np.unique options'''
u = np.unique(a[:,0])
idx = np.flatnonzero(np.diff(a[:,0])) + 1
return _custom_return(u, a, idx, return_groups)


def numpy_groupby_bins(a, return_groups=False):
'''Replace np.unique by np.bincount'''
bins = np.bincount(a[:,0])
nonzero_bins_idx = bins != 0
nonzero_bins = bins[nonzero_bins_idx]
idx = np.cumsum(nonzero_bins[:-1])
return _custom_return(np.flatnonzero(nonzero_bins_idx), a, idx, return_groups)


def numba_groupby_bins(a, return_groups=False):
'''Replace np.bincount by numba_bincount'''
bins = numba_bincount(a[:,0])
nonzero_bins_idx = bins != 0
nonzero_bins = bins[nonzero_bins_idx]
idx = np.cumsum(nonzero_bins[:-1])
return _custom_return(np.flatnonzero(nonzero_bins_idx), a, idx, return_groups)

所以 numba_bincount的工作方式和 np.bincount一样，它的定义是这样的:

from numba import njit


@njit
def _numba_bincount(a, counts, m):
for i in range(m):
counts[a[i]] += 1


def numba_bincount(arr): #just a refactor of Python count
M = np.max(arr)
counts = np.zeros(M + 1, dtype=int)
_numba_bincount(arr, counts, len(arr))
return counts

用法:

a = np.array([[1,275],[1,441],[1,494],[1,593],[2,679],[2,533],[2,686],[3,559],[3,219],[3,455],[4,605],[4,468],[4,692],[4,613]])
a = a[a[:, 0].argsort()]
>>> numpy_groupby_index(a, return_groups=False)
[array([275, 441, 494, 593]),
array([679, 533, 686]),
array([559, 219, 455]),
array([605, 468, 692, 613])]
>>> numpy_groupby_index(a, return_groups=True)
(array([1, 2, 3, 4]),
[array([275, 441, 494, 593]),
array([679, 533, 686]),
array([559, 219, 455]),
array([605, 468, 692, 613])])

性能测试

在我的计算机上对100M 个项目进行排序(使用10个不同的 ID)需要大约30秒。让我们测试剩余部分的方法运行所需的时间:

%matplotlib inline
benchit.setparams(rep=3)


sizes = [3*10**(i//2) if i%2 else 10**(i//2) for i in range(17)]
N = sizes[-1]
x1 = np.random.randint(0,10, size=N)
x2 = np.random.normal(loc=500, scale=200, size=N).astype(int)
a = np.transpose([x1, x2])


arr = a[a[:, 0].argsort()]
fns = [numpy_groupby_index, numpy_groupby_counts, numpy_groupby_diff, numpy_groupby_bins, numba_groupby_bins]
in_ = {s/1000000: (arr[:s], ) for s in sizes}
t = benchit.timings(fns, in_, multivar=True, input_name='Millions of events')
t.plot(logx=True, figsize=(12, 6), fontsize=14)

enter image description here

No doubt numba-powered bincount is a new winner of datasets that contains small IDs. It helps to reduce grouping of sorted data ~5 times which is ~10% of total runtime.

小开

阿什维尼 · 乔杜里建议的另一种方法可能正是您所要寻找的

def np_groupby(x, index):
return np.split(x, np.where(np.diff(x[:,index]))[0]+1)

X = numpy 数组

索引 = 列索引

[0] + 1 according to Ashwini, ... 任何非零的东西意味着它旁边的项目是不同的，我们可以使用 numpy.where来查找非零项目的索引，然后向它加1，因为这些项目的实际索引比返回的索引多一个; ... numpy.diff 用来查找项目实际上在哪里改变。


提交答案