Tri-Diagonal Matrix Algorithm （三重対角行列Aの連立一次方程式Ax=bの解のO(n)のアルゴリズム）

TDMA.py

# Tri-Diagonal Matrix Algorithm
def tdma(A, b):
  n = b.size
  p = []
  q = []

  p.append(A[0][1] / A[0][0])
  q.append(b[0] / A[0][0])  

  for i in range(1, n):
    if i != n-1:
      p.append(A[i][i+1] / (A[i][i] - A[i][i-1] * p[i-1]))

    q.append((b[i] - A[i][i-1] * q[i-1]) / (A[i][i] - A[i][i-1] * p[i-1]))


  ans = np.zeros(n)

  ans[n-1] = q[n-1]
  for i in range(n-2, -1, -1):
    ans[i] = q[i] - p[i] * ans[i+1]

  return ans

reference: https://sites.lsa.umich.edu/shixumeng/wp-content/uploads/sites/629/2020/02/Thomas.pdf