numpy: заполнить смещение по диагонали с различными значениями

Question

Мне нужно сделать матрицу n*nm, элементы которой следуют m(i,i+1)=sqrt(i) и 0 в противном случае. Например, для n=5, мы должны иметь

[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] 

где {a,b,c,d}=sqrt({1,2,3,4}). Here - решение для постоянной трехдиагональной матрицы, но мой случай немного сложнее. Я знаю, что могу сделать это с помощью цикла или со списком, но есть ли другие способы? n может быть потенциально большой.

например. (Список коды постижения)

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

Одним из способов может быть, чтобы создать массив из нулей, а затем использовать индексацию, чтобы выбрать и заполнить нужные индексы с Корневыми значениями.

Например:

>>> 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.  ]]) 

Answer 2

Вы можете использовать прямолинейно np.diag:

>>> d = np.sqrt(1 + np.arange(4)) 
>>> np.diag(d, 1) 
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.  ]]) 

Второй аргумент np.diag определяет диагональ в вопросе.