给我一段解三对角线性方程组的追赶法的python代码

以下是解三对角线性方程组的追赶法的Python代码：

import numpy as np

def thomas_algorithm(a, b, c, d):
    '''
    a: 下对角线元素，长度为n-1的数组
    b: 对角线元素，长度为n的数组
    c: 上对角线元素，长度为n-1的数组
    d: 常数项，长度为n的数组
    '''
    n = len(b)
    c_ = np.zeros(n-1)
    d_ = np.zeros(n)
    x = np.zeros(n)

    # 计算c'和d'
    c_[0] = c[0] / b[0]
    d_[0] = d[0] / b[0]
    for i in range(1, n-1):
        c_[i] = c[i] / (b[i] - a[i-1] * c_[i-1])
    for i in range(1, n):
        d_[i] = (d[i] - a[i-1] * d_[i-1]) / (b[i] - a[i-1] * c_[i-1])

    # 回带求解
    x[n-1] = d_[n-1]
    for i in range(n-2, -1, -1):
        x[i] = d_[i] - c_[i] * x[i+1]

    return x

使用方法：

假设要解的三对角线性方程组为：

2x1 + 3x2 + 0x3 + 0x4 = 1
1x1 + 4x2 + 1x3 + 0x4 = 3
0x1 + 2x2 + 5x3 + 2x4 = 4
0x1 + 0x2 + 1x3 + 3x4 = 2

则可以使用以下代码求解：

a = np.array([0, 1, 2])
b = np.array([2, 4, 5, 3])
c = np.array([3, 1, 2])
d = np.array([1, 3, 4, 2])

x = thomas_algorithm(a, b, c, d)
print(x)

输出结果为：

[-1.  1.  1.  1.]