如何根据列值对Numpy 2D矩阵中的行进行分组？

Question

什么是一种有效的(时间,简单)方法,2D NumPy通过不同的列条件(例如,按列2的值分组)对矩阵行进行分组,f1()并f2()在每个组上运行？

谢谢

Answer 1

如果你有一个arr形状数组(rows, cols),你可以得到第2列中所有值的向量

col = arr[:, 2]

然后,您可以使用分组条件构造一个布尔数组,例如,第1组由第2列中值大于5的那些行组成:

idx = col > 5

您可以将此布尔数组直接应用于原始数组以选择行:

group_1 = arr[idx]
group_2 = arr[~idx]

例如:

>>> arr = np.random.randint(10, size=(6,4))
>>> arr
array([[0, 8, 7, 4],
       [5, 2, 6, 9],
       [9, 5, 7, 5],
       [6, 9, 1, 5],
       [8, 0, 5, 8],
       [8, 2, 0, 6]])
>>> idx = arr[:, 2] > 5
>>> arr[idx]
array([[0, 8, 7, 4],
       [5, 2, 6, 9],
       [9, 5, 7, 5]])
>>> arr[~idx]
array([[6, 9, 1, 5],
       [8, 0, 5, 8],
       [8, 2, 0, 6]])

Answer 2

一个紧凑的解决方案是使用numpy_indexed(免责声明:我是它的作者),它实现了这种类型问题的完全矢量化解决方案:

最简单的使用方法是:

import numpy_indexed as npi
npi.group_by(arr[:, col1]).mean(arr)

但这也有效:

# run function f1 on each group, formed by keys which are the rows of arr[:, [col1, col2]
npi.group_by(arr[:, [col1, col2]], arr, f1)