演算法實現/線性代數/三對角矩陣演算法
所有提供的 三對角矩陣演算法 實現都假設三個對角線，a（下方），b（主對角線）和 c（上方），作為引數傳遞。
thomas :: Fractional g => [g] -> [g] -> [g] -> [g] -> [g]
thomas as bs cs ds = xs
where
n = length bs
bs' = b(0) : [b(i) - a(i)/b'(i-1) * c(i-1) | i <- [1..n-1]]
ds' = d(0) : [d(i) - a(i)/b'(i-1) * d'(i-1) | i <- [1..n-1]]
xs = reverse $ d'(n-1) / b'(n-1) : [(d'(i) - c(i) * x(i+1)) / b'(i) | i <- [n-2, n-3..0]]
-- convenience accessors (because otherwise it's hard to read)
a i = as !! (i-1) -- because the list's first item is equivalent to a_1
b i = bs !! i
c i = cs !! i
d i = ds !! i
x i = xs !! i
b' i = bs' !! i
d' i = ds' !! i
請注意，該函式接受一個子對角線列表（as），其第一個元素（索引 0）位於第 1 行，這就是為什麼我們在方便的訪問器中移動索引的原因，這樣我們就可以稱之為“a(1)”。
以下是在 C 程式語言中實現此演算法的示例。
void solve_tridiagonal_in_place_destructive(float * restrict const x, const size_t X, const float * restrict const a, const float * restrict const b, float * restrict const c) {
/*
solves Ax = v where A is a tridiagonal matrix consisting of vectors a, b, c
x - initially contains the input vector v, and returns the solution x. indexed from 0 to X - 1 inclusive
X - number of equations (length of vector x)
a - subdiagonal (means it is the diagonal below the main diagonal), indexed from 1 to X - 1 inclusive
b - the main diagonal, indexed from 0 to X - 1 inclusive
c - superdiagonal (means it is the diagonal above the main diagonal), indexed from 0 to X - 2 inclusive
Note: contents of input vector c will be modified, making this a one-time-use function (scratch space can be allocated instead for this purpose to make it reusable)
Note 2: We don't check for diagonal dominance, etc.; this is not guaranteed stable
*/
c[0] = c[0] / b[0];
x[0] = x[0] / b[0];
/* loop from 1 to X - 1 inclusive, performing the forward sweep */
for (size_t ix = 1; ix < X; ix++) {
const float m = 1.0f / (b[ix] - a[ix] * c[ix - 1]);
c[ix] = c[ix] * m;
x[ix] = (x[ix] - a[ix] * x[ix - 1]) * m;
}
/* loop from X - 2 to 0 inclusive (safely testing loop condition for an unsigned integer), to perform the back substitution */
for (size_t ix = X - 1; ix-- > 0; )
x[ix] -= c[ix] * x[ix + 1];
x[0] -= c[0] * x[1];
}
以下變體保留了方程組以供在其他輸入上重用。注意庫呼叫分配和釋放暫存空間的必要性 - 在許多輸入上解決相同的三對角線系統的更有效實現將依賴於呼叫函式提供指向暫存空間的指標。
void solve_tridiagonal_in_place_reusable(float * restrict const x, const size_t X, const float * restrict const a, const float * restrict const b, const float * restrict const c) {
/* Allocate scratch space. */
float * restrict const cprime = malloc(sizeof(float) * X);
if (!cprime) {
/* do something to handle error */
}
cprime[0] = c[0] / b[0];
x[0] = x[0] / b[0];
/* loop from 1 to X - 1 inclusive */
for (size_t ix = 1; ix < X; ix++) {
const float m = 1.0f / (b[ix] - a[ix] * cprime[ix - 1]);
cprime[ix] = c[ix] * m;
x[ix] = (x[ix] - a[ix] * x[ix - 1]) * m;
}
/* loop from X - 2 to 0 inclusive, safely testing loop end condition */
for (size_t ix = X - 1; ix-- > 0; )
x[ix] -= cprime[ix] * x[ix + 1];
x[0] -= cprime[0] * x[1];
/* free scratch space */
free(cprime);
}
以下 C++ 函式將解決一個一般的三對角線系統（儘管它會在過程中破壞輸入向量 c 和 d）。注意索引 這裡是從 零開始 的，換句話說 其中 是未知數的個數。注意，除了在 main() 函式中列印文字之外，此程式碼在 C 和 C++ 中都是有效的。
#include <iostream>
using namespace std;
void solve(double* a, double* b, double* c, double* d, int n) {
/*
// n is the number of unknowns
|b0 c0 0 ||x0| |d0|
|a1 b1 c1||x1|=|d1|
|0 a2 b2||x2| |d2|
1st iteration: b0x0 + c0x1 = d0 -> x0 + (c0/b0)x1 = d0/b0 ->
x0 + g0x1 = r0 where g0 = c0/b0 , r0 = d0/b0
2nd iteration: | a1x0 + b1x1 + c1x2 = d1
from 1st it.: -| a1x0 + a1g0x1 = a1r0
-----------------------------
(b1 - a1g0)x1 + c1x2 = d1 - a1r0
x1 + g1x2 = r1 where g1=c1/(b1 - a1g0) , r1 = (d1 - a1r0)/(b1 - a1g0)
3rd iteration: | a2x1 + b2x2 = d2
from 2nd it. : -| a2x1 + a2g1x2 = a2r2
-----------------------
(b2 - a2g1)x2 = d2 - a2r2
x2 = r2 where r2 = (d2 - a2r2)/(b2 - a2g1)
Finally we have a triangular matrix:
|1 g0 0 ||x0| |r0|
|0 1 g1||x1|=|r1|
|0 0 1 ||x2| |r2|
Condition: ||bi|| > ||ai|| + ||ci||
in this version the c matrix reused instead of g
and the d matrix reused instead of r and x matrices to report results
Written by Keivan Moradi, 2014
*/
n--; // since we start from x0 (not x1)
c[0] /= b[0];
d[0] /= b[0];
for (int i = 1; i < n; i++) {
c[i] /= b[i] - a[i]*c[i-1];
d[i] = (d[i] - a[i]*d[i-1]) / (b[i] - a[i]*c[i-1]);
}
d[n] = (d[n] - a[n]*d[n-1]) / (b[n] - a[n]*c[n-1]);
for (int i = n; i-- > 0;) {
d[i] -= c[i]*d[i+1];
}
}
int main() {
int n = 4;
double a[4] = { 0, -1, -1, -1 };
double b[4] = { 4, 4, 4, 4 };
double c[4] = {-1, -1, -1, 0 };
double d[4] = { 5, 5, 10, 23 };
// results { 2, 3, 5, 7 }
solve(a,b,c,d,n);
for (int i = 0; i < n; i++) {
cout << d[i] << endl;
}
cout << endl << "n= " << n << endl << "n is not changed hooray !!";
return 0;
}
Python
[edit | edit source]注意索引 這裡是從 零開始 的，換句話說 其中 是未知數的個數。
try:
import numpypy as np # for compatibility with numpy in pypy
except:
import numpy as np # if using numpy in cpython
def TDMASolve(a, b, c, d):
n = len(a)
ac, bc, cc, dc = map(np.array, (a, b, c, d))
xc = []
for j in range(1, n):
if(bc[j - 1] == 0):
ier = 1
return
ac[j] = ac[j]/bc[j-1]
bc[j] = bc[j] - ac[j]*cc[j-1]
if(b[n-1] == 0):
ier = 1
return
for j in range(1, n):
dc[j] = dc[j] - ac[j]*dc[j-1]
dc[n-1] = dc[n-1]/bc[n-1]
for j in range(n-2, -1, -1):
dc[j] = (dc[j] - cc[j]*dc[j+1])/bc[j]
return dc
另一種方法
def TDMA(a,b,c,f):
a, b, c, f = map(lambda k_list: map(float, k_list), (a, b, c, f))
alpha = [0]
beta = [0]
n = len(f)
x = [0] * n
for i in range(n-1):
alpha.append(-b[i]/(a[i]*alpha[i] + c[i]))
beta.append((f[i] - a[i]*beta[i])/(a[i]*alpha[i] + c[i]))
x[n-1] = (f[n-1] - a[n-2]*beta[n-1])/(c[n-1] + a[n-2]*alpha[n-1])
for i in reversed(range(n-1)):
x[i] = alpha[i+1]*x[i+1] + beta[i+1]
return x
# small test. x = (1,2,3)
if __name__ == '__main__':
a = [3,3]
b = [2,1]
c = [6,5,8]
f = [10,16,30]
print(TDMA(a,b,c,f))
MATLAB
[edit | edit source]waitfunction x = TDMAsolver(a,b,c,d)
%a, b, c are the column vectors for the compressed tridiagonal matrix, d is the right vector
n = length(d); % n is the number of rows
% Modify the first-row coefficients
c(1) = c(1) / b(1); % Division by zero risk.
d(1) = d(1) / b(1);
for i = 2:n-1
temp = b(i) - a(i) * c(i-1);
c(i) = c(i) / temp;
d(i) = (d(i) - a(i) * d(i-1))/temp;
end
d(n) = (d(n) - a(n) * d(n-1))/( b(n) - a(n) * c(n-1));
% Now back substitute.
x(n) = d(n);
for i = n-1:-1:1
x(i) = d(i) - c(i) * x(i + 1);
end
Fortran 90
[edit | edit source]注意索引 這裡是從一開始的，換句話說 其中 是未知數的個數。
有時不希望求解器例程覆蓋三對角線係數（例如，對於求解多個方程組，其中只有方程組的右側發生變化），因此這種實現給出了一個相對便宜的方法來儲存係數。
subroutine solve_tridiag(a,b,c,d,x,n)
implicit none
! a - sub-diagonal (means it is the diagonal below the main diagonal)
! b - the main diagonal
! c - sup-diagonal (means it is the diagonal above the main diagonal)
! d - right part
! x - the answer
! n - number of equations
integer,parameter :: r8 = kind(1.d0)
integer,intent(in) :: n
real(r8),dimension(n),intent(in) :: a,b,c,d
real(r8),dimension(n),intent(out) :: x
real(r8),dimension(n) :: cp,dp
real(r8) :: m
integer i
! initialize c-prime and d-prime
cp(1) = c(1)/b(1)
dp(1) = d(1)/b(1)
! solve for vectors c-prime and d-prime
do i = 2,n
m = b(i)-cp(i-1)*a(i)
cp(i) = c(i)/m
dp(i) = (d(i)-dp(i-1)*a(i))/m
end do
! initialize x
x(n) = dp(n)
! solve for x from the vectors c-prime and d-prime
do i = n-1, 1, -1
x(i) = dp(i)-cp(i)*x(i+1)
end do
end subroutine solve_tridiag
此子例程提供覆蓋 d 或不覆蓋 d 的選項。 [1]
subroutine tdma(n,a,b,c,d,x)
implicit none
integer, intent(in) :: n
real, intent(in) :: a(n), c(n)
real, intent(inout), dimension(n) :: b, d
real, intent(out) :: x(n)
! --- Local variables ---
integer :: i
real :: q
! --- Elimination ---
do i = 2,n
q = a(i)/b(i - 1)
b(i) = b(i) - c(i - 1)*q
d(i) = d(i) - d(i - 1)*q
end do
! --- Backsubstitution ---
q = d(n)/b(n)
x(n) = q
do i = n - 1,1,-1
q = (d(i) - c(i)*q)/b(i)
x(i) = q
end do
return
end