Extracting classes from numpy array

Question

I am having a numpy array of data

data = np.random.random((5, 5))

and another numpy array of equal shape masking the first one with classes from 0 to n :

>>> mask
array([[3, 3, 1, 1, 0],
       [2, 0, 1, 2, 2],
       [0, 1, 0, 0, 3],
       [2, 1, 1, 0, 2],
       [0, 2, 3, 0, 2]])

What's the best way to compute a two-dimensional array with n rows, where each row contains all elements from data with class row_idx (described by mask )?

Answer 1

You can't do it better than O(n^2) as you should iterate through all mask array. As the number of elements in each class can be different the result can't be numpy array (rows have different sizes). So I don't think you can avoid python loop here with pure numpy functions. I suggest this O(n^2) solution:

ans = [[] for i in range(mask.max() + 1)]
for k, v in np.ndenumerate(mask):
    ans[v].append(data[k])