Home > Back-end >  Replacing diagonal elements of matrix in Python
Replacing diagonal elements of matrix in Python

Time:04-11

The code counts the number of ones in each row. Then I want it to replace the diagonal with this count. However, it takes only the last count. The current and desired outputs are attached.

import numpy as np
num_ones=np.zeros(1)
y = np.array([[0, 0, 0], 
             [1, 0, 1],
             [1, 0, 0]])
for x in range(0, 3):
    num_ones = (y[x] == 1).sum()
    print([num_ones])

np.fill_diagonal(y, [-num_ones])
print([y])

Current output:

[array([[-1,  0,  0],
       [ 1, -1,  1],
       [ 1,  0, -1]])]

Desired output:

[array([[0,  0,  0],
       [ 1, -2,  1],
       [ 1,  0, -1]])]

CodePudding user response:

Don't loop, use vectorial code to sum:

import numpy as np
num_ones=np.zeros(1)
y = np.array([[0, 0, 0], 
              [1, 0, 1],
              [1, 0, 0]])

d = -(y==1).sum(1)
# or, if only 0/1
# d = -y.sum(1)

np.fill_diagonal(y, d)

print(y)

Output:

[[ 0  0  0]
 [ 1 -2  1]
 [ 1  0 -1]]
  • Related