【連立一次方程式】TDMA（三重対角行列アルゴリズム、Tri-Diagonal Matrix Algorithm）C++コード付き

今回は連立一次方程式の係数が三重対角行列のときにとても効率の良い、TDMA（三重対角行列アルゴリズム、Tri-Diagonal Matrix Algorithm）のC++コードを公開します。TDMAは直接法に分類される手法で、流体計算の際のソルバーとして用いられています。理論は

http://penguinitis.g1.xrea.com/study/note/tdma.pdf

がわかりやすく参考になります。方程式が3つの例がわざわざ示されており、とてもわかりやすいです。アルゴリズムの流れを書きます。連立一次方程式を

$\left( \begin{array}{ccc} a_1 & b_1 & & & & 0 \\ c_2 & a_2 & b_2 & & & & \\ & c_3 & a_3 & b_3 & & & \\ & & \ddots & \ddots & \ddots & \\ & & & c_{n-1} & a_{n-1} & b_{n-1} \\ 0 & & & & a_{n} & b_{n} \\ \end{array} \right) \left( \begin{array}{c} x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_{n-1} \\ x_n \end{array} \right) = \left( \begin{array}{c} d_1 \\ d_2 \\ d_3 \\ \vdots \\ d_{n-1} \\ d_n \end{array} \right)$

とします。このときTDMAのアルゴリズムは

1. $P_1=-\frac{b_1}{a_1}$ 、 $Q_1=\frac{d_1}{a_1}$ と定める
2. $P_i=\frac{-b_i}{a_i+c_i P_{i-1}}$ 、 $Q_i=\frac{d_i-c_i Q_{i-1}}{a_i+c_i P_{i-1}}$ を使って $i=2$ から $i=n$ までの $P_i$ と $Q_i$ を求める
3. $x_n=Q_n$ とする
4. $x_i=P_ix_{i+1}+Q_i$ を使って $i=n-1$ から $i=1$ までの $x_i$ を求める（後ろ向きであることに注意！）

となります。

今回解く連立一次方程式は

$Ax=d$

とあらわされます。ここで、 $A$ は $4 \times 4$ の行列、 $x$ と $d$ は $4$ 次元のベクトルです。 $A$ と $d$ が既知で、 $x$ が未知です。具体的には

$\left( \begin{array}{ccc} 2 & 1 & 0 & 0 \\ 1 & 2 & 1 & 0 \\ 0 & 1 & 2 & 1 \\ 0 & 0 & 1 & 2 \end{array} \right) \left( \begin{array}{c} x_1 \\ x_2 \\ x_3 \\ x_4 \end{array} \right) = \left( \begin{array}{c} 4 \\ 8 \\ 12 \\ 11 \end{array} \right)$

であり、これを解くと $x_1=1, x_2=2, x_3=3, x_4=4$ となります。

コード上ではそれぞれ対角要素を配列 t[i] に、対角右の要素を配列 r[i] に、対角左の要素を配列 l[i] に収納しています。ではコードを以下に示します。

#include <iostream>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <vector>

using namespace std;

//TDMA//
inline void tdma(double a[], double c[], double b[], double d[], double x[], int size)
{
	int i;
	double *P=new double[size];
	double *Q=new double[size];
	
	//first step//
	P[0]=-b[0]/a[0];
	Q[0]=d[0]/a[0];
	
	//second step//
	for(i=1;i<size;i++)
	{
		P[i]=-b[i]/(a[i]+c[i]*P[i-1]);
		Q[i]=(d[i]-c[i]*Q[i-1])/(a[i]+c[i]*P[i-1]);
	}
	
	//third step//
	x[size-1]=Q[size-1];

	//fourth step, backward//
	for(i=size-2;i>-1;i=i-1)
	{
		x[i]=P[i]*x[i+1]+Q[i];
	}
	
	delete [] P,Q;
}

int main()
{
	int i;
	int size=4;
	
	double *t=new double[size];
	double *l=new double[size];
	double *r=new double[size];
	double *b=new double[size];
	double *x=new double[size];
	
	t[0]=2.0;t[1]=2.0;t[2]=2.0;t[3]=2.0;
	l[0]=0.0;l[1]=1.0;l[2]=1.0;l[3]=1.0;
	r[0]=1.0;r[1]=1.0;r[2]=1.0;r[3]=0.0;
	b[0]=4.0;b[1]=8.0;b[2]=12.0;b[3]=11.0;
	
	
	tdma(t,l,r,b,x,size);
	
	for(i=0;i<size;i++) cout<<x[i]<<endl;
	
	delete [] t,l,r,b,x;

    return 0;
}