python - 获取与标签对应的行，对于许多标签

Question

我有一个二维数组，其中每一行都有一个存储在单独数组中的标签(不一定是唯一的)。对于每个标签，我想从我的二维数组中提取具有此标签的行。我想要的一个基本工作示例是:

import numpy as np

data=np.array([[1,2],[3,5],[7,10], [20,32],[0,0]])
label=np.array([1,1,1,0,1])

#very simple approach
label_values=np.unique(label)
res=[]
for la in label_values:
    data_of_this_label_val=data[label==la]
    res+=[data_of_this_label_val]
print(res)

结果 (res) 可以是任何格式，只要它易于访问即可。在上面的例子中，它将是

[array([[20, 32]]), array([[ 1,  2],
   [ 3,  5],
   [ 7, 10],
   [ 0,  0]])]

请注意，我可以轻松地将列表中的每个元素与 label_values 中的唯一标签之一相关联(即按索引)。

虽然这可行，但使用 for 循环可能会花费大量时间，尤其是当我的标签向量很大时。这可以加快速度或更优雅地编码吗？

Answer 1

您可以argsort 标签(我相信这就是 unique 的作用)。

如果您的标签是示例中的小非负整数，您可以获得更便宜的价格，请参阅 https://stackoverflow.com/a/53002966/7207392 .

>>> import numpy as np
>>> 
>>> data=np.array([[1,2],[3,5],[7,10], [20,32],[0,0]])
>>> label=np.array([1,1,1,0,1])
>>> 
>>> idx = label.argsort()
# use kind='mergesort' if you require a stable sort, i.e. one that
# preserves the order of equal labels
>>> ls = label[idx]
>>> split = 1 + np.where(ls[1:] != ls[:-1])[0]
>>> np.split(data[idx], split)
[array([[20, 32]]), array([[ 1,  2],
       [ 3,  5],
       [ 7, 10],
       [ 0,  0]])]