python - 如何根据内容从 numpy 数组中提取行？

Question

例如，作为标题，我有一个 2d numpy 数组，如下所示，

[[33, 21, 1],
 [33, 21, 2],
 [32, 22, 0],
 [33, 21, 3],
 [34, 34, 1]]

我想根据第一列和第二列的内容依次提取这些行，在这种情况下，我想得到 3 个不同的 2d numpy 数组，如下所示，

[[33, 21, 1],
 [33, 21, 2],
 [33, 21, 3]]

和

[[32, 22, 0]]

和

[[34, 34, 1]]

我可以使用 numpy 中的什么函数来执行此操作？我认为重点是用第一列和第二列来区分不同的行。如果这些列中的元素相同，则特定行将归类在同一输出数组中。 我想写一个 python 函数来完成这种工作，因为我可以有一个比上面的大得多的数组。请随时给我建议，谢谢。

Answer 1

这是处理许多此类分组的方法 -

# Sort array based on second column
sorted_a = a[np.argsort(a[:,1])]

# Get shifting indices for first col. Split along axis=0 using those.
shift_idx = np.unique(sorted_a[:,1],return_index=True)[1][1:]
out = np.split(sorted_a,shift_idx)

或者，为了提高性能，我们可以获得 shift_idx，就像这样 -

shift_idx = np.flatnonzero(sorted_a[1:,1] > sorted_a[:-1,1])+1

sample 运行-

In [27]: a
Out[27]: 
array([[33, 21,  1],
       [33, 21,  2],
       [32, 22,  0],
       [33, 21,  3],
       [34, 34,  1]])
In [28]: sorted_a = a[np.argsort(a[:,1])]

In [29]: np.split(sorted_a,np.unique(sorted_a[:,1],return_index=True)[1][1:])
Out[29]: 
[array([[33, 21,  1],
        [33, 21,  2],
        [33, 21,  3]]), array([[32, 22,  0]]), array([[34, 34,  1]])]