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