1

I want to set the smallest k values in each rows to 0,without using for loop. Here is my Code with for loop:

import numpy as np
k = 2
sims = np.array([[3,1,2,9],[5,9,1,7],[1,8,6,2], [1,5,8,9]])
for i in range(len(sims)):
    indices_argsort = np.argsort(sims[i]) 
    sims[i, indices_argsort[: -k]] = 0
print(sims)

The output is :

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

In the output , the smallest k values can be set to 0. But forloop is bad to use when deal with large matrices. So, is there some solutions can solve without using forloop?

MarioKZZ
  • 141
  • 2
  • 8

2 Answers2

0

I have solve it with Broadcast mechanism as follows:

import numpy as np
sims = np.array([[3,1,2,9],[5,9,1,7],[1,8,6,2], [1,5,8,9]])
indices_argsort = np.argsort(sims , axis = -1)

k = 2

sims[np.arange(4)[:,None], indices_argsort[:,:-k]] = 0
print(sims)

output:

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

Use

mask = np.apply_along_axis(lambda x: x < np.partition(x, k)[k], 1, sims)
sims[mask] = 0

Output

array([[3, 0, 0, 9],
       [0, 9, 0, 7],
       [0, 8, 6, 0],
       [0, 0, 8, 9]])
Vishnudev Krishnadas
  • 10,679
  • 2
  • 23
  • 55