如何将Numpy数组每一行的最小k值设置为0？

Question

我想将每行中最小的 k 值设置为 0，而不使用 for 循环。 这是我的带有 for 循环的代码：

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)

输出是：

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

在输出中，可以将最小的 k 值设置为 0。但是在处理大型矩阵时，使用 forloop 是不好的。那么，有没有不使用forloop也能解决的方案呢？

Answer 1

使用

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

输出

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

Answer 2

我用广播机制解决了如下：

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)

输出：

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