追赶法求解三对角矩阵(C/C++/python/matlab/Fortran 90)

原文地址
Python地址

追赶法求解过程
import numpy as np

## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver
def TDMAsolver(a, b, c, d):
    '''
    TDMA solver, a b c d can be NumPy array type or Python list type.
    refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
    and to http://www.cfd-online.com/Wiki/Tridiagonal_matrix_algorithm_-_TDMA_(Thomas_algorithm)
    '''
    nf = len(d) # number of equations
    ac, bc, cc, dc = map(np.array, (a, b, c, d)) # copy arrays
    for it in range(1, nf):
        mc = ac[it-1]/bc[it-1]
        bc[it] = bc[it] - mc*cc[it-1] 
        dc[it] = dc[it] - mc*dc[it-1]
                
    xc = bc
    xc[-1] = dc[-1]/bc[-1]

    for il in range(nf-2, -1, -1):
        xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]

    return xc

例如:

A = np.array([[10,2,0,0],[3,10,4,0],[0,1,7,5],[0,0,3,4]],dtype=float)   

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

print(TDMAsolver(a, b, c, d))
>> [ 0.14877589  0.75612053 -1.00188324  2.25141243]
#compare against numpy linear algebra library
print(np.linalg.solve(A, d))
>> [ 0.14877589  0.75612053 -1.00188324  2.25141243]
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容

  • 戴安娜去世二十年了。时间冲刷着英伦玫瑰的存在感,办公室的小鲜肉竟不知道她了。 她那么的美丽、优雅、能干。 她那么地...
    汝婴阅读 426评论 0 0
  • 能量指数:6 自己的人生已经按下重启键,从硬件到软件,到各个程序,全部进行了重装最新最优系统,并且进行全方位的升级...
    神乐酱酱酱阅读 6,812评论 0 1
  • 有一次参加一个游戏的可用性测试,主持人让我完成一些任务,但在他们表示我已经完成了之后,我自己还在继续寻找操作的路径...
    类猿人阅读 351评论 0 1