灵感来自Eelco Hoogendoorn 图书馆 https://stackoverflow.com/a/38015063/1488055,但没有他的库,并利用数组的第一列总是增加的事实(如果不是,请首先使用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])]
我没有“timeit”([编辑]见下文),但这可能是解决问题的更快方法:
- 没有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 个。
我对熊猫感到惊讶,我跑了好几次,但它在这种情况下表现得很糟糕。
>>> 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 的回答 https://stackoverflow.com/a/53859634/1488055 and 贝赫扎德·沙耶格 https://stackoverflow.com/users/13146793/behzad-shayegh(参见评论)
还要感谢TMBailey https://stackoverflow.com/users/16327476/tmbailey注意 argsort 的复杂性是 n log(n)。