Selecting rows from numpy ndarray when groups are known

Question

I have a numpy array x and a list of tags cls which preserves the order of x, and records the class each element from x belongs to. For example, for two different classes 0 and 1:

x = [47 21 16 19 38]
cls = [0 0 1 0 1]

meaning that 47, 21, 19 belong together, and so do 16, 38.

Is there a pythonic way to select elements from x by class?

Now I'm doing

for clusterId in np.unique(cls):
indices = [i for i in range(len(cls)) if cls[i]==clusterId]
print 'Class ', clusterId
for idx in indices:
    print '\t', x[idx,].tolist()

but I can't believe there is nothing more elegant.

Answer 1

cls is perfect for constructing a boolean index array:

>>> import numpy as np
>>> x = np.array([47, 21, 16, 19, 38])
>>> cls = np.array([0, 0, 1, 0, 1],dtype=bool)
>>> x[cls]
array([16, 38])
>>> x[~cls]
array([47, 21, 19])

Note that if cls isn't a boolean array already, you can make it one using ndarray.astype:

>>> cls = np.array([0, 0, 1, 0, 1])
>>> x[cls]  #Not what you want
array([47, 47, 21, 47, 21])
>>> x[cls.astype(bool)]  #what you want.
array([16, 38])

In the general case where you have a cls array and you want to pick out only elements with that index:

>>> x[cls == 0]
array([47, 21, 19])
>>> x[cls == 1]
array([16, 38])

To find all the classes that you have, you can use np.unique(cls)