Python - How to replace elements other than principal diagonal and anti-diagonal to zero in a matrix taking n * n matrix as input in python

Question

example taking input : 5 5 taking matrix input is :

1 2 3 4 1 
5 6 8 1 5 
1 5 6 4 4
8 2 6 2 6
1 5 6 8 9

Expected output :

1 0 0 0 1
0 6 0 1 0
0 0 6 0 0
0 2 0 2 0
1 0 0 0 9

enter code here
m,n= input().split()
new_matrix = []
for i in range(int(m)):
    a = list(map(int,input().split()))
    new_matrix.append(a);

Answer 1

To replace all elements other than at diagonal/anti-diagonal in n*n matrix you can do:

matrix = [
    [1, 2, 3, 4, 1],
    [5, 6, 8, 1, 5],
    [1, 5, 6, 4, 4],
    [8, 2, 6, 2, 6],
    [1, 5, 6, 8, 9],
]

n = len(matrix)

for i in range(n):
    for j in range(n):
        if i != j and j != n - i - 1:
            matrix[i][j] = 0

print(matrix)

Prints:

[
    [1, 0, 0, 0, 1],
    [0, 6, 0, 1, 0],
    [0, 0, 6, 0, 0],
    [0, 2, 0, 2, 0],
    [1, 0, 0, 0, 9],
]