Keep max N values per row in a Numpy array

Question

I need to keep the max N (3) values per row in an Array.

a=np.array([[1,2,3,4],[8,7,6,5],[5,3,1,2]])
a
Out[135]: 
array([[1, 2, 3, 4],
       [8, 7, 6, 5],
       [5, 3, 1, 2]])

The indexes of those can be identified with np.partition:

n=3
np.argpartition(a, -n, axis=1)[:,-n:]
Out[136]: 
array([[1, 2, 3],
       [2, 1, 0],
       [3, 0, 1]], dtype=int64)

So, my question is: How should I keep values from those indices and set to zero others to get:

Out[136]: 
array([[0, 2, 3, 4],
       [8, 7, 6, 0],
       [5, 3, 0, 2]])

Jelle Westra · Accepted Answer · 2021-07-07 17:44:55Z

2

a=np.array([[1,2,3,4],[8,7,6,5],[5,3,1,2]])

n=3
mask = np.argpartition(a, -n, axis=1) < a.shape[1] - n

a[mask] = 0

answered Jul 7, 2021 at 17:44

Jelle Westra

5463 silver badges8 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Andras Deak -- Слава Україні · Accepted Answer · 2021-07-07 17:46:36Z

0

One option is to use your indices as fancy indices to overwrite the values in a zero array:

import numpy as np

# input
a = np.array([[1,2,3,4], [8,7,6,5], [5,3,1,2]])
n = 3

indices = (np.arange(a.shape[0]), np.argpartition(a, -n, axis=1)[:,-n:])
res = np.zeros_like(a)
res[indices] = a[indices]

Then you get

>>> res
array([[0, 2, 3, 4],
       [8, 7, 6, 0],
       [5, 3, 0, 2]])

edited Jul 7, 2021 at 17:46

answered Jul 7, 2021 at 17:41

Andras Deak -- Слава Україні

35.4k13 gold badges94 silver badges118 bronze badges

Collectives™ on Stack Overflow

Keep max N values per row in a Numpy array

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related