小能豆

有没有numpy group by函数?

javascript

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)

阅读 38

收藏
2024-07-12

共1个答案

小能豆

要按第一列对 NumPy 数组进行分组,可以使用numpy函数以及itertools.groupby列表推导。以下是实现此目的的分步方法:

  1. 按第一列对数组进行排序。
  2. 用于itertools.groupby按第一列对已排序的数组进行分组。
  3. 将分组的值收集到一个新的数组中。

代码如下:

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])]

解释:

  1. 排序:使用 根据第一列对数组进行排序a[a[:, 0].argsort()]
  2. 分组itertools.groupby用于按第一列对元素进行分组。
  3. 收集结果:我们将分组的元素收集到列表列表中,其中每个子列表包含特定组的第二列值。
  4. 数组转换:将结果转换为 NumPy 对象数组,以方便使用。

此代码应实现按第一列对数组进行所需的分组。

2024-07-12