Python/Numpy: get top k largest values in a 2D matrix as a mask

user_11235205 0
Viewed 265

Let's say I have a 3x3 matrix like this:

array([[8, 6, 3],
       [6, 7, 2],
       [0, 8, 9]])

Now I want to get the top k largest values in the matrix, and create a mask from it. If the number is in the top k largest, it has value 1, else 0. Let k=2. In the example above there are one 9 and two 8, we need to take all of them, so the returned mask is like this:

array([[1, 0, 0],
       [0, 0, 0],
       [0, 1, 1]])

I have read this and that answer, and I can use the indices as the mask. However, I wonder if there is any better solution?

numpy arrays python
1 Answers

How about this?

def is_topk(a, k=1):
    _, rix = np.unique(-a, return_inverse=True)
    return np.where(rix < k, 1, 0).reshape(a.shape)

Example on your array:

>>> is_topk(a, 1)
array([[0, 0, 0],
       [0, 0, 0],
       [0, 0, 1]])

>>> is_topk(a, 2)
array([[1, 0, 0],
       [0, 0, 0],
       [0, 1, 1]])
user_7581740
Related