Setting increasing values in a numpy array around a defined diagonal

Question

What is the best way to create a 2D list (or numpy array) in python, in which the diagonal is set to -1 and the remaining values are increasing from 0 by 1, for different values of n. For example, if n = 3 the array would look like:

[[-1,0,1]
 [2,-1,3]
 [4,5,-1]] 

or for n = 4:

[[-1,0,1,2]
 [3,-1,4,5]
 [6,7,-1,8]
 [9,10,11,-1]]

etc.

I know I can create an array with zeros and with the diagonal set to -1 with: a = numpy.zeros((n,n)) numpy.fill_diagonal(a,-1)

And so if n = 3 this would give:

[[-1,0,0]
 [0,-1,0]
 [0,0,-1]]

But how would I then set the 0's to be increasing numbers, as shown in the example above? Would I need to iterate through and set the values through a loop? Or is there a better way to approach this?

Thanks in advance.

Answer 1

One approach -

def set_matrix(n):
    out = np.full((n,n),-1)
    off_diag_mask = ~np.eye(n,dtype=bool)
    out[off_diag_mask] = np.arange(n*n-n)
    return out

Sample runs -

In [23]: set_matrix(3)
Out[23]: 
array([[-1,  0,  1],
       [ 2, -1,  3],
       [ 4,  5, -1]])

In [24]: set_matrix(4)
Out[24]: 
array([[-1,  0,  1,  2],
       [ 3, -1,  4,  5],
       [ 6,  7, -1,  8],
       [ 9, 10, 11, -1]])