0

I'm looking for the quickest way to do the following:

arr = np.array([2.,0.5,1.3,4.5,7.4])

I want to build a new array (or overwrite current) where I put to zero everything except the n highest entries.

[2.0,0,0,4.5,7.4]
gsamaras
  • 71,951
  • 46
  • 188
  • 305
James Arten
  • 523
  • 5
  • 16
  • Is it possible for `arr` to have duplicate entries? If so, what output do you expect when, say `arr = [1, 2, 2, 3, 3, 3]` with `n = 2`? – Warren Weckesser Feb 07 '22 at 17:45
  • 1
    Does this answer your question? [How do I get indices of N maximum values in a NumPy array?](https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array) – Agaz Wani Feb 07 '22 at 17:47

3 Answers3

2

You can use argsort and indexing. This solution is in place:

arr = np.array([2.,0.5,1.3,4.5,7.4])

N = 3

# get the sorted order
idx = np.argsort(arr)[:-N]

# replace all but the N highest with 0
arr[idx]= 0

output: array([2. , 0. , 0. , 4.5, 7.4])

mozway
  • 194,879
  • 13
  • 39
  • 75
0

I am not sure whether it is the fastest or not but it gives the expected output

import numpy as np
array = np.array([2.,0.5,1.3,4.5,7.4])
l = len(array)
order = array.argsort()
ranks = order.argsort()

n = 3
array = [array[i] if ranks[i] >= l - n else 0 for i in range(len(ranks))]
print(array)
ksohan
  • 1,165
  • 2
  • 9
  • 23
0

Let's define the array and a value for N, so we keep the N highest elements.

N = 3
arr = np.array([2.,0.5,1.3,4.5,7.4])

Using argsort we can retrieve the order of elements from lower to highest. We just need to subtract from the size of the array to get from highest to lower.

order = arr.argsort()
order = arr.size - order

Then replace with zeros the smaller elements:

arr[order > N] = 0
Walisson Cardoso
  • 109
  • 3