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

数学作业要求实现共轭梯度法的算法。
题目中的矩阵A是n=400/500/600的三对角矩阵。
在网上查阅资料未果后，自己解决了。

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)

结果输出：

C:\Users\87167\Anaconda3\envs\tf\python.exe C:/Users/87167/PycharmProjects/GongETiDuFa/main.py
[[-2  0  0  0  0]
 [ 0 -2  0  0  0]
 [ 0  0 -2  0  0]
 [ 0  0  0 -2  0]
 [ 0  0  0  0 -2]]
[[0 0 0 0 0]
 [1 0 0 0 0]
 [0 1 0 0 0]
 [0 0 1 0 0]
 [0 0 0 1 0]]
[[0 1 0 0 0]
 [0 0 1 0 0]
 [0 0 0 1 0]
 [0 0 0 0 1]
 [0 0 0 0 0]]
[[-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]]