Построение диагональных и недиагональных матричных элементов

Question

Код Mathematica

В Mathematica мне удалось выписать нужную матрицу с диагональными и недиагональными значениями. Мне было интересно, как лучше всего это сделать в python используя numpy?

Answer 1

Виртуальный клон вашего кода:

In [146]: arr = np.zeros((5,5),int)                                                            
In [147]: arr[np.arange(5),np.arange(5)]=2                                                     
In [148]: arr[np.arange(4),np.arange(1,5)]=-1                                                  
In [149]: arr[np.arange(1,5),np.arange(4)]=-1                                                  
In [150]: arr                                                                                  
Out[150]: 
array([[ 2, -1,  0,  0,  0],
       [-1,  2, -1,  0,  0],
       [ 0, -1,  2, -1,  0],
       [ 0,  0, -1,  2, -1],
       [ 0,  0,  0, -1,  2]])

или с функцией diag:

In [151]: np.diag(np.ones(5,int)*2,0)+np.diag(np.ones(4,int)*-1,-1)+np.diag(np.ones(4,int)*-1,1
     ...: )                                                                                    
Out[151]: 
array([[ 2, -1,  0,  0,  0],
       [-1,  2, -1,  0,  0],
       [ 0, -1,  2, -1,  0],
       [ 0,  0, -1,  2, -1],
       [ 0,  0,  0, -1,  2]])