Given a numpy array of the form below:

x = [[4.,3.,2.,1.,8.],[1.2,3.1,0.,9.2,5.5],[0.2,7.0,4.4,0.2,1.3]]

is there a way to retain the top-3 values in each row and set others to zero in python (without an explicit loop). The result in the case of the example above would be

x = [[4.,3.,0.,0.,8.],[0.,3.1,0.,9.2,5.5],[0.0,7.0,4.4,0.0,1.3]]

Code for one example

import numpy as np
arr = np.array([1.2,3.1,0.,9.2,5.5,3.2])
indexes=arr.argsort()[-3:][::-1]
a = list(range(6))
A=set(indexes); B=set(a)
zero_ind=(B.difference(A)) 
arr[list(zero_ind)]=0

The output:

array([0. , 0. , 0. , 9.2, 5.5, 3.2])

Above is my sample code (with many lines) for a 1-D numpy array. Looping through each row of a numpy array and performing this same computation repeatedly would be quite expensive. Is there a simpler way?

Here is a fully vectorized code without third party outside numpy. It is using numpy's argpartition to efficiently find the k-th values. See for instance this answer for other use cases.

def truncate_top_k(x, k, inplace=False):m, n = x.shape# get (unsorted) indices of top-k valuestopk_indices = numpy.argpartition(x, -k, axis=1)[:, -k:]# get k-th valuerows, _ = numpy.indices((m, k))kth_vals = x[rows, topk_indices].min(axis=1)# get boolean mask of values smaller than k-this_smaller_than_kth = x < kth_vals[:, None]# replace mask by 0if not inplace:return numpy.where(is_smaller_than_kth, 0, x)x[is_smaller_than_kth] = 0return x