python - 在numpy数组中将每行的K个随机非零元素设置为零，矢量化解决方案

Question

我想以矢量化的方式将 numpy 数组的每一行的 K 个随机非零元素设置为零。

我的非矢量化代码如下所示，带有典型的输入和输出。

谢谢并恭祝安康，

伊格尔

orig_arr = np.array([[1, 1, 0, 1], [0, 0, 1, 1], [0, 1, 1, 1], [1, 1, 0, 0], [1, 1, 1, 1]])
k = 2
for i in range(len(orig_arr)):
    this_row_elements = np.nonzero(orig_arr[i, :])[0]
    sampled = random.sample(list(this_row_elements), k)
    orig_arr[i, sampled] = 0

Before:
    [[1 1 0 1]
     [0 0 1 1]
     [0 1 1 1]
     [1 1 0 0]
     [1 1 1 1]]
After:
    [[0 1 0 0]
     [0 0 0 0]
     [0 0 1 0]
     [0 0 0 0]
     [0 1 0 1]]

标签： pythonarraysnumpyrandomvectorization