numpy: fill offset diagonal with different values

Question

I need to make a n*n matrix m whose elements follow m(i,i+1)=sqrt(i) and 0 otherwise. For example, for n=5, we should have

[0 a 0 0 0]
[0 0 b 0 0]
[0 0 0 c 0]
[0 0 0 0 d]
[0 0 0 0 0]

where {a,b,c,d}=sqrt({1,2,3,4}). Here is a solution for a constant tri-diagonal matrix, but my case is a bit more complicated than that. I know I can do that with a loop or with list comprehension, but are there other ways? n can be potentially big.

e.g. (list comprehension code)

ele=lambda i,j:sqrt(i+1) if j-i==1 else 0

[[ele(i,j) for j in range(0,6)] for i in range(0,6)]

Answer 1

One way could be to create the array of zeros and then use indexing to select and fill the desired indices with the square-root values.

For example:

>>> z = np.zeros((5,5))
>>> rng = np.arange(4)
>>> z[rng, rng+1] = np.sqrt(rng+1)
>>> z
array([[ 0.        ,  1.        ,  0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  1.41421356,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  1.73205081,  0.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  2.        ],
       [ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ]])