numpy中是否有函數替換numpy數組的低對角線值和高對角線值？

Question

要更換主對角線我用np.fill_diagonal：

matrix = np.zeros((4, 4), float) 
main = np.array([2,2,2,2]) 
np.fill_diagonal(matrix, main) 

，但我也需要更換上下對角線旁邊，是主對角線：

upper=np.array([1,1,1]) 
lower=np.array([7,7,7]) 

獲得：

matrix=[[2 1 0 0] 
     [7 2 1 0] 
     [0 7 2 1] 
     [0 0 7 2]] 

謝謝

Answer 1

隨着一些智能切片，np.fill_diagonal也可以這樣做：

>>> np.fill_diagonal(matrix[:-1, 1:], upper) 
>>> np.fill_diagonal(matrix[1:, :-1], lower) 
>>> matrix 
array([[ 2., 1., 0., 0.], 
     [ 7., 2., 1., 0.], 
     [ 0., 7., 2., 1.], 
     [ 0., 0., 7., 2.]])