Python求解三对角方程组LU分解

原文是matlab版的，我改写了一个Python版本的

(4条消息) (追赶法求三对角矩阵、LU分解)_三对角矩阵的lu分解_Giannis_34的博客-CSDN博客

#!/usr/bin/python
import numpy as np

#solves the system of linear equations Ax = res
#where A is a tridiagonal matrix with diag2 representing the main diagonal
#diag1 and diag3 represent the lower and upper diagonal respectively
#all four parameters are vectors of size n
#the first element of diag1 and the last element of diag3 are ignored
#therefore diag1[i], diag2[i] and diag3[i] are located on the same row of A
def solve_tridiagonal_equation(diag1, diag2, diag3, res):
    #assert(len(diag1) == len(diag2) == len(diag3) == len(res))
    n=len(res)
    y=np.zeros(n)
    solution=np.zeros(n)
    beta=np.ones(n)
    v=np.ones(n-1)
    beta[0]=diag2[0]
    v[0]=diag3[0]/diag2[0]
    for i in range(1,n-1):
        beta[i]=diag2[i]-diag1[i]*v[i-1]
        v[i]=diag3[i]/beta[i]
    
    beta[n-1]=diag2[n-1]-diag1[n-1]*v[n-2]
    
    y[0]=res[0]/diag2[0]
    for i in range(1,n):
        y[i]=(res[i]-diag1[i]*y[i-1])/beta[i]
        
    solution[n-1]=y[n-1]
    for i in range(n-2,-1,-1):
        solution[i]=y[i]-v[i]*solution[i+1]
    
    return solution

测试一下

a=[0,1,1,1,1]
b=[2,2,2,2,2]
c=[1,1,1,1,0]
d=[3,4,4,4,3]
solve_tridiagonal_equation(a,b,c,d)

结果是 

[1., 1., 1., 1., 1.]