python - 如果超过变化的行索引，则将 numpy 矩阵元素设置为零

Question

我有一个相当大的 m次n numpy 矩阵 M填充非零值和数组 x长度为 m，其中每个条目表示行索引，之后矩阵元素应设置为零。例如，如果 n=5和 x[i]=3 , 然后矩阵的第 i 行设置为 [M_i1, M_i2, M_i3, 0, 0] .

如果 x 的所有条目具有相同的值 k , 我可以简单地使用像 M[:,k:]=0 这样的切片, 但我想不出一种有效的方法来为每一行使用不同的值而不遍历所有行并对每一行使用切片。

我想创建一个看起来像 [[1]*x[1] + [0]*(n-x[1]),...,[1]*x[m] + [0]*(n-x[m])] 的矩阵并将它用于 bool 索引，但也不知道如何在不循环的情况下创建它。

非矢量化解决方案如下所示:

for i in range(m):
    if x[i] < n:
        M[i,x[i]:] = 0

示例输入

M = np.array([[1,2,3],[4,5,6]])
m, n =  2, 3
x = np.array([1,2])

和输出

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

有人有这个问题的矢量化解决方案吗？

非常感谢!

Answer 1

您可以使用多维 bool 索引:

M[x[:,None]<=np.arange(M.shape[1])] = 0

例子:

M = [[7, 8, 4, 2, 3, 9, 1, 8, 4, 3],
     [2, 1, 6, 1, 5, 2, 2, 2, 9, 2],
     [6, 1, 6, 8, 4, 3, 6, 9, 2, 6],
     [5, 4, 0, 8, 3, 0, 0, 1, 8, 7],
     [8, 7, 8, 8, 9, 2, 0, 8, 0, 2]]
x = [4, 4, 0, 6, 2]

输出:

      [[7, 8, 4, 2, 0, 0, 0, 0, 0, 0],
       [2, 1, 6, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [5, 4, 0, 8, 3, 0, 0, 0, 0, 0],
       [8, 7, 0, 0, 0, 0, 0, 0, 0, 0]]