如何根据条件和位置构造矩阵？

Question

我有一个矩阵A =

np.matrix([[0, 0, 0, 0, 0],
        [0, 0, 0, 1, 1],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 1, 1, 1]]

如果A [i，j] = 1和（i + 1）％3 = 0，我想建立一个矩阵B，其中B [i，j] = 5；否则B [i，j] = 0。

B应该是：B =

np.matrix([[0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 5, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 5, 0, 0]]

是否有任何可能的方法可以不使用for循环来实现此目的，就像矩阵计算一样？谢谢。

Answer 1

更新后的答案：

我还没有想过消除列表理解中的for循环以进行剩余条件过滤的方法，但是计算的“繁重”部分利用了numpy的优化。

import numpy as np
newdata = 5 * np.array([(i + 1) % 3 == 0 for i in range(data.shape[-1])]) * np.array(data)

原始答案（在不能使用for循环的条件下）：

假设矩阵存储为data，则可以使用列表推导语法来获取所需的内容。

newdata = [[5 if val == 1 and (idx + 1) % 3 == 0 else 0
            for idx, val in enumerate(row)]
           for row in data]