How to retain N smallest elements in a given row of numpy array?

user_8893595 0
Viewed 461

Given a 2-D numpy matrix, how to retain the N smallest elements in each row and change rest of them to 0(zero).

For example: N=3 Input array:

1   2   3   4   5
4   3   6   1   0
6   5   3   1   2

Expected output:

1   2   3   0   0
0   3   0   1   0
0   0   3   1   2

Following is the code that I have tried and it works:

# distance_matrix is the given 2D array
N=3
for i in range(distance_matrix.shape[0]):
    n_th_largest = np.sort(distance_matrix[i])[N]
    for j in range(distance_matrix.shape[1]):
        distance_matrix[i][j] = np.where(distance_matrix[i][j]<n_th_largest,distance_matrix[i][j],0)

# return distance_matrix

However, this operation involves iterating over every single element. Is there a faster way to solve this using np.argsort() or any other function?

import numpy sorting python
2 Answers

Approach #1

Here's one with np.argpartition for performance efficiency -

N = 3
newval = 0
np.put_along_axis(a,np.argpartition(a,N,axis=1)[:,N:],newval,axis=1)

Explanation : We partition the input array to get indices that are partitioned-across for the kth argument in np.argpartition. So, basically consider this as two partitions, with first one for smallest N elements along that axis and the other for the rest. We need to reset the second partition, which we select with [:,N:] and we use np.put_along_axis to do the resetting.

Sample run -

In [144]: a # input array
Out[144]: 
array([[1, 2, 3, 4, 5],
       [4, 3, 6, 1, 0],
       [6, 5, 3, 1, 2]])

In [145]: np.put_along_axis(a,np.argpartition(a,3,axis=1)[:,3:],0,axis=1)

In [146]: a
Out[146]: 
array([[1, 2, 3, 0, 0],
       [0, 3, 0, 1, 0],
       [0, 0, 3, 1, 2]])

Approach #2

Here's another again with np.argpartition, but just slicing the Nth smallest element per row and then resetting all greater than it. As such, if there are duplicates for the Nth smallest element, we will keep all those with this method. Here's the implementation -

a[a>=a[np.arange(len(a)), np.argpartition(a,3,axis=1)[:,3],None]] = 0

Timings on a scaled up version -

In [184]: a = np.array([[1,2,3,4,5],[4,3,6,1,0],[6,5,3,1,2]])

In [185]: a = np.repeat(a,10000,axis=0)

In [186]: %timeit np.put_along_axis(a,np.argpartition(a,3,axis=1)[:,3:],0,axis=1)
1.78 ms ± 5.89 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [187]: a = np.array([[1,2,3,4,5],[4,3,6,1,0],[6,5,3,1,2]])

In [188]: a = np.repeat(a,10000,axis=0)

In [189]: %timeit a[a>=a[np.arange(len(a)), np.argpartition(a,3,axis=1)[:,3],None]] = 0
1.54 ms ± 54.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
user_32938810

If we can use pandas and have your input in a dataframe, this is a one-liner with .apply(..., axis=1) on each row:

df.apply(lambda row: row.nsmallest(3), axis=1).fillna(0).astype(int)

   0  1  2  3  4
0  1  2  3  0  0
1  0  3  0  1  0
2  0  0  3  1  2

Notes:

And here's the boilerplate to make your example reproducible:

import pandas as pd
from io import StringIO

dat = """1   2   3   4   5
4   3   6   1   0
6   5   3   1   2"""

df = pd.read_csv(StringIO(dat), sep='\s+', header=None)
user_2022290
Related