如何根据内容从numpy数组中提取行? [英] How to extract rows from an numpy array based on the content?

问题描述

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

As title, for example, I have an 2d numpy array, like the one below,

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

并且我想根据第一列和第二列中的内容有序地提取这些行，在这种情况下，我想获得3个不同的2d numpy数组，如下所示，

and I want to extract these rows orderly based on the content in the first and the second column, in this case, I want to get 3 different 2d numpy arrays, as below,

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

和

[[32, 22, 0]]

和

[[34, 34, 1]]

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

What function in numpy could I use to do this? I think the point is to distinguish different rows with their first and second columns. If elements in these columns are the same, then the specific rows are categorized in the same output array. I want to write a python function to do this kind of job, because I could have a much more bigger array than the one above. Feel free to give me advice, thank you.

推荐答案

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

Here's an approach to handle many such groupings -

# 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-

Alternatively, for performance efficiency purposes, we can get shift_idx, like so -

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

样品运行-

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

这篇关于如何根据内容从numpy数组中提取行?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！