Replace the top 10 values in numpy

Question

Is there any easy way to replace the top 10 values with 1 and the rest of them with zeros? I have found that numpy argpartition can give me a new array with the index but I haven't been able to easily use it in the original array? Can anyone help? Thanks in Advance

Answer 1

You could do it using np.sort to find the 10th largest value, and then use np.where to flag the array.

import numpy as np

a = np.random.rand(30)

a_10 = np.sort(a)[-10]

a_new = np.where(a >= a_10, 1, 0)

print(a)     # Print the original
print(a_new) # Print the boolean array

EDIT: A single-line, in-place operation is thus

a = np.where(a >= np.sort(a)[-10], 1, 0)

EDIT2: The answer can be extended to 2D. I made a 6x6 matrix, where I flag per row the 3 largest values with a 1.

# 2D example, save top3 per 
a = np.random.rand(6, 6)

a_3 = np.sort(a, axis=1)[:,-3]
a_new = np.where(a >= a_3[:,None], 1, 0)

print(a)
print(a_new)