算法实现/线性代数/三对角矩阵算法
所有提供的 三对角矩阵算法 实现都假设三个对角线，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