Python生成三对角矩阵

import numpy as np
def ThreediaMatrix(down,mid,up,n):         
    array_a=np.diag([down]*(n-1),-1)    # down是对角线下一行的元素，末尾是-1，n-1是元素的个数
    array_b = np.diag([mid] * n)        # mid是对角线上的元素，默认为0
    array_c=np.diag([up]*(n-1),1)       # up是对角线上面的元素，末尾是1
    matrix_A=array_a+array_b+array_c
    return  matrix_A
A=ThreediaMatrix(-1,2,1,5)
print(A)

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