CODEWITHSUNDEEP
pythonnumpy
James Arten
James Arten

Reputation: 666

Put everything to zero in a np.array except n highest entries

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]

Upvotes: 0

Views: 496

Answers (3)

Walisson Cardoso
Walisson Cardoso

Reputation: 109

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

Upvotes: 0

mozway
mozway

Reputation: 261015

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])

Upvotes: 2

ksohan
ksohan

Reputation: 1203

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)

Upvotes: 0

Related Questions