高效地将零分配给二维numpy数组中每一行的多列

时间:2018-08-29 12:22:49

标签: numpy matrix-indexing numpy-indexing

我想为2d numpy矩阵分配零,其中对于每一行,我们都有一个截止列,之后应将其设置为零。

例如,在这里我们有尺寸为4x5的矩阵A,带有截止列[1,3,2,4]。 我想做什么:

import numpy as np
np.random.seed(1)
A = np.random.rand(4, 5)
cutoff = np.array([1,3,2,4])
A[0, cutoff[0]:] = 0
A[1, cutoff[1]:] = 0
A[2, cutoff[2]:] = 0
A[3, cutoff[3]:] = 0

我可以使用np.repeat行索引和列来做到这一点,但是我的矩阵太大了,我无法做到这一点。有没有一种有效的方法可以做到这一点?

1 个答案:

答案 0 :(得分:0)

使用broadcasting创建完整的掩码并分配-

A[cutoff[:,None] <= np.arange(A.shape[1])] = 0

或者使用builtin外部方法-

A[np.less_equal.outer(cutoff, range(A.shape[1]))] = 0