Subset a np.matrix efficiently based on boolean np.array but only until certain treshold

user_10743830 0
Viewed 168

So lets say we have a Matrix M

M=np.array([[1,2,3],
            [1,2,3],
            [1,2,3],
            [1,2,3]
            [1,2,3],
            [1,2,3]])

with the same number of rows as the length of a np.array mask:

mask = np.array([False,True,False,True,False,True])

And there is additional Parameter called threshold=2

I want to subset M with mask only until 2 True values, for all of the remaining I would set false. That means that M[mask] should return me second and fourth row only not the last one. Is there an efficient way to do this with numpy avoiding for loops?

numpy-ndarray numpy python
2 Answers

M[mask][:2]

this selects your rows according to the mask and stops after reaching 2 True values.

user_78138470

You can go along two ways:

import numpy as np

M = np.array([[1,2,3],
              [1,2,3],
              [1,2,3],
              [1,2,3],
              [1,2,3],
              [1,2,3]])

mask = np.array([False,True,False,True,False,True])

true_locs = np.where(mask)[0]

# set True to False in the mask
mask[true_locs[2:]] = False

# OR just use the indeces directly
M[true_locs[:2],:]
user_129981470
Related