追赶法求解Ax=b python语言示例

下面是使用追赶法求解Ax=b的Python语言示例：

import numpy as np

def thomas_algorithm(A, b):
    # 获取矩阵A的行数
    n = len(A)
    # 初始化追赶法需要的变量
    a = np.zeros(n-1)
    c = np.zeros(n-1)
    d = np.zeros(n)
    x = np.zeros(n)
    
    # 将矩阵A转化为追赶法需要的三对角矩阵形式
    a[0] = A[0][1] / A[0][0]
    c[n-2] = A[n-2][n-1] / A[n-2][n-2]
    for i in range(1, n-1):
        a[i] = A[i][i+1] / (A[i][i] - A[i][i-1]*a[i-1])
    for i in range(n-3, -1, -1):
        c[i] = A[i][i+1] / (A[i][i] - A[i][i+1]*c[i+1])
    
    # 前向追赶
    d[0] = b[0] / A[0][0]
    for i in range(1, n):
        d[i] = (b[i] - A[i][i-1]*d[i-1]) / (A[i][i] - A[i][i-1]*a[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

其中，A是一个n×n的矩阵，b是一个n维的向量，返回的x是一个n维的向量，表示Ax=b的解。