有没有numpy group by函数？

问题描述：

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)

解决方案 1：

受到Eelco Hoogendoorn 的图书馆的启发，但没有他的图书馆，并且利用了数组第一列始终在增加的事实（如果不是，请先用排序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:])