Here's one vectorized approach with np.argpartition
-
m,n = a.shape
a[np.arange(m)[:,None],np.argpartition(a,n-2,axis=1)[:,:-2]] = 0
Sample run -
In [570]: a
Out[570]:
array([[ 0.94791114, 0.48438182, 0.54574317, 0.45481231, 0.94013836],
[ 0.03861196, 0.99047316, 0.7897759 , 0.38863967, 0.93659426],
[ 0.49436676, 0.93762758, 0.33694977, 0.45701655, 0.73078113],
[ 0.21240062, 0.85141765, 0.00815352, 0.52517721, 0.49752736]])
In [571]: m,n = a.shape
...: a[np.arange(m)[:,None],np.argpartition(a,n-2,axis=1)[:,:-2]] = 0
...:
In [572]: a
Out[572]:
array([[ 0.94791114, 0. , 0. , 0. , 0.94013836],
[ 0. , 0.99047316, 0. , 0. , 0.93659426],
[ 0. , 0.93762758, 0. , 0. , 0.73078113],
[ 0. , 0.85141765, 0. , 0.52517721, 0. ]])