Getting rows corresponding to label, for many labels

user_8921867 0
Viewed 1832

I have a 2D array, where each row has a label that is stored in a separate array (not necessarily unique). For each label, I want to extract the rows from my 2D array that have this label. A basic working example of what I want would be this:

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)

The result (res) can have any format, as long as it is easily accessible. In the above example, it would be

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

Note that I can easily associate each element in my list to one of the unique labels in label_values (that is, by index).

While this works, using a for loop can take quite a lot of time, especially if my label vector is large. Can this be sped up or coded more elegantly?

python-3.x numpy sorting arrays python
2 Answers

You can argsort the labels (which is what unique does under the hood I believe).

If your labels are small nonnegatvie integers as in the example you can get it a bit cheaper, see 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]])]
user_72073920

Unfortunately, there isn't a built-in groupby function in numpy, though you could write alternatives. However, your problem could be solved more succinctly using pandas, if that's available to you:

import pandas as pd

res = pd.DataFrame(data).groupby(label).apply(lambda x: x.values).tolist()
# or, if performance is important, the following will be faster on large arrays, 
# but less readable IMO:
res = [data[i] for i in pd.DataFrame(data).groupby(label).groups.values()]

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