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)
要按第一列对 NumPy 数组进行分组,可以使用numpy函数以及itertools.groupby列表推导。以下是实现此目的的分步方法:
numpy
itertools.groupby
代码如下:
import numpy as np from itertools import groupby # Sample data 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]]) # Step 1: Sort the array by the first column a_sorted = a[a[:, 0].argsort()] # Step 2: Group by the first column using groupby grouped = groupby(a_sorted, key=lambda x: x[0]) # Step 3: Collect the results into a list of arrays result = np.array([list(map(lambda x: x[1], group)) for key, group in grouped], dtype=object) # Print the result print(result)
输出:
[array([275, 441, 494, 593]) array([679, 533, 686]) array([559, 219, 455]) array([605, 468, 692, 613])]
a[a[:, 0].argsort()]
此代码应实现按第一列对数组进行所需的分组。