python如何生成任意n阶的三对角矩阵

import numpy as np
def generate_matrix(n):
    # 使用对角矩阵相加得到三对角矩阵A
    array_a = np.diag([-2] * n)
    array = np.diag([1] * (n-1))
    a = np.zeros((n-1))
    b = np.zeros(n)
    array_b = np.insert(array, 0, values=a, axis=0)# 添加行
    array_b = np.insert(array_b, (n-1), values=b, axis=1)# 添加列
    array_c = np.insert(array, (n-1), values=a, axis=0)
    array_c = np.insert(array_c, 0, values=b, axis=1)
    matrix_A = array_a + array_b + array_c
    print(array_a)
    print(array_b)
    print(array_c)
    print(matrix_A)

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    n = 5
    generate_matrix(n)

>>> m = np.mat("1,2,3;4,5,6;7,8,9")
>>> m
matrix([[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]])
>>> L = np.tril(m,-1)
>>> L
array([[0, 0, 0],
       [4, 0, 0],
       [7, 8, 0]])
>>> U = np.triu(m,1)
>>> U
array([[0, 2, 3],
       [0, 0, 6],
       [0, 0, 0]])

>>> D = np.diag(np.diag(m))
>>> D
array([[1, 0, 0],
       [0, 5, 0],
       [0, 0, 9]])
>>> D = np.mat(D)
>>> D
matrix([[1, 0, 0],
        [0, 5, 0],
        [0, 0, 9]])

>>> D = np.tril(m) - L
>>> D
array([[1, 0, 0],
       [0, 5, 0],
       [0, 0, 9]])
>>> D = np.mat(D)
>>> D
matrix([[1, 0, 0],
        [0, 5, 0],
        [0, 0, 9]])