根据更改阈值过滤np数组中值的最快方法

Question

I want to filter an array arr based on some thresholds.

arr = np.array([2,2,2,2,2,5,5,5,1])
thresholds = np.array([4,1])

I want to filter arr based on the values in thresholds when the value in arr is greater than the threshold

My idea is to create a mask for each threshold

Expected result:

# [[False False False False False  True  True  True False]
#  [ True  True  True  True  True  True  True  True False]]

One way to do it in Python:

mask = [True if x>condi else False for condi in thresholds for x in arr]
mask = np.reshape(mask,(2,9))

Then to get the filtered array by just filteredarr = arr[mask[i]] where i is the index of the relevant threshold

Is there a better way (performance wise) to do it in Python ? Especially that I am dealing with big arrays (len around 250000 for arr, no specific len for thresholds yet, but I am expecting a big array) ?

Edit: The final output expected on the data is [array([5, 5, 5]), array([2, 2, 2, 2, 2, 5, 5, 5])]

Answer 1

The mask can easily be obtained using

mask = arr[None,:]>thresholds[:,None]
mask

# Output
# array([[False, False, False, False, False,  True,  True,  True, False],
#        [ True,  True,  True,  True,  True,  True,  True,  True, False]], dtype=bool)

The idea is to blow up the dimensionality by adding an additional axis using None (which does the same as np.newaxis ) and to compare then the arrays element-wise.

Once we have the mask we can filter the data using various methods where the choice strongly depends on your problem:

• Of course you can do

   res = [arr[m] for m in mask] # [array([5, 5, 5]), array([2, 2, 2, 2, 2, 5, 5, 5])] 

  in order to obtain a list with the filtered data, but it is slow in general.

• In case you have further numeric calculations I would create a masked array in which only the filtered data are taken into account:

   m = np.zeros_like(mask).astype(np.int) m[:] = arr res = np.ma.masked_where(~mask,m) 

  Each line corresponds now to the filtered data according to the corresponding threshold. Masked arrays allow you to continue working with many functions like mean or std

   res.mean(axis=1) # masked_array(data = [5.0 3.125], # mask = [False False], # fill_value = 1e+20) res.mean(axis=1).compressed() # array([ 5. , 3.125])