如何減少2D numpy代碼中if語句的數量？

Question

我有一個numpy ndarray（ M ），其中包含一個包含N x N個數字（N是整數）的2D對稱矩陣。

[[1. 2. 0. 0. 2. 2. 2.]
 [2. 1. 2. 0. 0. 2. 2.]
 [0. 2. 1. 2. 0. 0. 2.]
 [0. 0. 2. 1. 2. 0. 0.]
 [2. 0. 0. 2. 1. 2. 0.]
 [2. 2. 0. 0. 2. 1. 2.]
 [2. 2. 2. 0. 0. 2. 1.]]

我想檢查一下：將M中的項目i，j設置為0，將i + 1，i + 2，i + 3中的j設置為M中的零

我有這段代碼，這樣做：

for i in range(len(M)):
     for j in range(len(M)):
        if M[i, j] != 0:
            if j == len(M) - 3:
                M[i, j+2] = 0
                M[i, j+1] = 0
            elif j == len(M) - 2:
                M[i, j+1] = 0
            elif j == len(M) - 1:
                   continue
            else:
                M[i, j + 3] = 0
                M[i, j + 2] = 0
                M[i, j + 1] = 0
        else:
            continue

它工作得很好，但是我想減少我的if語句。 我讀到有關M[i,j:j+3] ，但我真的不知道如何使用它。

我正在等待以下輸出矩陣：

[[1. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]
 [0. 2. 0. 0. 0. 0. 2.]
 [0. 0. 2. 0. 0. 0. 0.]
 [2. 0. 0. 0. 1. 0. 0.]
 [2. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]]

Answer 1

這應該工作：

for i in range(len(M)):
     for j in range(len(M)):
        if M[i, j] != 0:
            M[i, (j+1):(j+4)] = 0
print (M)

輸出：

[[1. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]
 [0. 2. 0. 0. 0. 0. 2.]
 [0. 0. 2. 0. 0. 0. 0.]
 [2. 0. 0. 0. 1. 0. 0.]
 [2. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]]