python - 是否有任何 numpy 按功能分组？

Question

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)

Answer 1

灵感来自 Eelco Hoogendoorn's library ，但没有他的库，并且使用数组的第一列总是增加的事实(如果没有，首先使用 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 上运行了 timeit，以获得 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 个。我对 pandas 感到惊讶，我跑了好几次，但它在这种情况下表现不佳。

>>> 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's answer和 Behzad Shayegh (引用评论)也感谢 TMBailey注意 argsort 的复杂度是 n log(n)。