How to replace the N smallest elements in each row of numpy array?-CodePudding

I would like to replace the N smallest elements in each row for 0, and that the resulting array would respect the same order and shape of the original array.

Specifically, if the original numpy array is:

 import numpy as np

 x = np.array([[0,50,20],[2,0,10],[1,1,0]])

And N = 2, I would like for the result to be the following:

x = np.array([[0,50,0],[0,0,10],[0,1,0]])

I tried the following, but in the last row it replaces 3 elements instead of 2 (because it replaces both 1s and not only one)

import numpy as np

N = 2
x = np.array([[0,50,20],[2,0,10],[1,1,0]])

x_sorted = np.sort(x , axis = 1)
x_sorted[:,max_edges:] = 0

replace = x_sorted.copy()

final = np.where(np.isin(x,replace),0,x)

Note that this is small example and I would like that it works for a much bigger matrix.

Thanks for your time!

CodePudding user response：

Use numpy.argpartition to find the index of N smallest elements, and then use the index to replace values:

N = 2
idy = np.argpartition(x, N, axis=1)[:, :N]
x[np.arange(len(x))[:,None], idy] = 0

x

array([[ 0, 50,  0],
       [ 0,  0, 10],
       [ 1,  0,  0]])

Notice if there are ties, it could be undetermined which values get replaced depending on the algorithm used.

CodePudding user response：

One way using numpy.argsort:

N = 2
x[x.argsort().argsort() < N] = 0

Output:

array([[ 0, 50,  0],
       [ 0,  0, 10],
       [ 0,  1,  0]])