Home > Software design >  A better way to create an NxN matrix, with specific diagonal and off diagonal elements with numpy
A better way to create an NxN matrix, with specific diagonal and off diagonal elements with numpy

Time:03-12

I need to make a NxN matrix and specify diagonal elements. This is what I tried so far, I was hoping to find a more elegant solution without loops.

N = 4
value_offdiag  = 2    

b = np.eye(N, N)
b[np.triu_indices(N, 1)] = value_offdiag
b[np.tril_indices(4, -1)] = value_offdiag

This works, but there's got to be a better way without using a loop. Just wanted to check (google search so far hasn't revealed much)

CodePudding user response:

What about using numpy.fill_diagonal?

N = 4
value_offdiag = 2
b = np.ones((N,N))*value_offdiag
np.fill_diagonal(b,1)

print(b)

output:

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

CodePudding user response:

One-liner:

N = 4
value1 = 100
value2 = 234

a = np.eye(N)*value1   abs(np.eye(N)-1)*value2

Output:

>>> a.astype(int)
array([[ 13, 240, 240, 240],
       [240,  13, 240, 240],
       [240, 240,  13, 240],
       [240, 240, 240,  13]])
  • Related