Take the N first values ​​from each row in the NumPy matrix that satisfy the condition

Question

I have a numpy vector and a numpy array .

I need to take from each row in the matrix the first N (say 3) values ​​that are less than (or equal to) the corresponding row in the vector.

so if this is my vector:

 7, 9, 22, 38, 6, 15 

and this is my matrix:

 [[ 20., 9., 7., 5., None, None], [ 33., 21., 18., 9., 8., 7.], [ 31., 21., 13., 12., 4., 0.], [ 36., 18., 11., 7., 7., 2.], [ 20., 14., 10., 6., 6., 3.], [ 14., 14., 13., 11., 5., 5.]] 

the conclusion should be:

 [[7,5,None], [9,8,7], [21,13,12], [36,18,11], [6,6,3], 14,14,13]] 

Is there an efficient way to do this with masks or something else without an ugly for loop?

Any help would be appreciated!

Answer 1

Approach No. 1

Here is one with broadcasting -

 def takeN_le_per_row_broadcasting(a, b, N=3): # a, b : 1D, 2D arrays respectively # First col indices in each row of b with <= corresponding one in a idx = (b <= a[:,None]).argmax(1) # Get all N ranged column indices all_idx = idx[:,None] + np.arange(N) # Finally advanced-index with those indices into b for desired output return b[np.arange(len(all_idx))[:,None], all_idx] 

Approach # 2

Inspired by NumPy Fancy Indexing - Crop different ROIs from different channels solution , we can use np.lib.stride_tricks.as_strided to efficiently retrieve patches, for example:

 from skimage.util.shape import view_as_windows def takeN_le_per_row_strides(a, b, N=3): # a, b : 1D, 2D arrays respectively # First col indices in each row of b with <= corresponding one in a idx = (b <= a[:,None]).argmax(1) # Get 1D sliding windows for each element off data w = view_as_windows(b, (1,N))[:,:,0] # Use fancy/advanced indexing to select the required ones return w[np.arange(len(idx)), idx]