Как я могу уменьшить количество операторов if в моем двумерном коде?

Question

У меня есть nundy ndarray (M), содержащий двумерную симметричную матрицу, которая содержит N x N чисел (N - целое число), например,

[[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.]]

Я хочу проверить это: установить элементi, j в M к 0, где j в i + 1, i + 2, i + 3 к нулю в 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.]]