Thomas algorithm
(Math/thomas.hpp)

• View this file on GitHub
• Last update: 2025-02-01 01:02:42+09:00
• Include: #include "Math/thomas.hpp"

概要

Tridiagonal_matrix_algorithm を参考。
三重対角行列についての解を $ \text{O}(N) $ で求める。

$ n $ 個の未知数を持つ三重対角方程式系は、次のように表すことができる。

\[a_{i}x_{i-1} + b_{i}x_{i} + c_{i}x_{i+1} = d_{i},\] \[a_{1} = 0 \quad \text{and} \quad c_{n} = 0.\]

として

\[\begin{bmatrix} b_1 & c_1 & 0 & & \\ a_2 & b_2 & c_2 & & \\ & a_3 & b_3 & \ddots & \\ & & \ddots & \ddots & c_{n-1} \\ 0 & & & a_n & b_n \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \\ x_3 \\ \vdots \\ x_n \end{bmatrix} = \begin{bmatrix} d_1 \\ d_2 \\ d_3 \\ \vdots \\ d_n \end{bmatrix}\]

例題

Codeforces Beta Round 24 D. Broken robot

Code

template<class T> vector<T> thomas(const vector<T> &a, const vector<T> &b, 
                                   const vector<T> &c, const vector<T> &d){
    const int n = b.size();
    vector<T> scratch(n), x(n);
    scratch[0] = c[0] / b[0];
    x[0] = d[0] / b[0];
    for (int i = 1; i < n; i++) {
        const T div = b[i] - a[i] * scratch[i - 1];
        if (i + 1 < n) scratch[i] = c[i] / div;
        x[i] = (d[i] - a[i] * x[i - 1]) / div;
    }

    for (int i = n - 2; i >= 0; i--)
        x[i] -= scratch[i] * x[i + 1];
    
    return x;
}

#line 1 "Math/thomas.hpp"
template<class T> vector<T> thomas(const vector<T> &a, const vector<T> &b, 
                                   const vector<T> &c, const vector<T> &d){
    const int n = b.size();
    vector<T> scratch(n), x(n);
    scratch[0] = c[0] / b[0];
    x[0] = d[0] / b[0];
    for (int i = 1; i < n; i++) {
        const T div = b[i] - a[i] * scratch[i - 1];
        if (i + 1 < n) scratch[i] = c[i] / div;
        x[i] = (d[i] - a[i] * x[i - 1]) / div;
    }

    for (int i = n - 2; i >= 0; i--)
        x[i] -= scratch[i] * x[i + 1];
    
    return x;
}

Back to top page