numpy 数组抽取_如何根据内容从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]]

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

and

[[32, 22, 0]]

and

[[34, 34, 1]]

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)

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

Sample run -

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