Прогонка для трёхдиагональных систем

Урок про метод прогонки (алгоритм Томаса) — специализированный Гаусс для трёхдиагональных систем, работающий за O(n).

Метод прогонки (алгоритм Томаса) решает систему, у которой ненулевые элементы только на главной и двух соседних диагоналях, за линейное время прямым и обратным проходом.

Где берутся трёхдиагональные системы

Это один из самых частых видов СЛАУ в численных методах. Они возникают при построении кубических сплайнов, при решении краевых задач для дифференциальных уравнений конечными разностями, в неявных схемах для уравнения теплопроводности. В таких системах каждое уравнение связывает неизвестное только с двумя соседями — отсюда три диагонали. Применять к ним общий Гаусс ($O(n^3)$) — расточительство: 99% элементов матрицы нули, и они остаются нулями в процессе. Прогонка использует эту структуру и решает систему за $O(n)$.

| b0 c0  0  0 |   | x0 |   | d0 |
| a1 b1 c1  0 | · | x1 | = | d1 |
|  0 a2 b2 c2 |   | x2 |   | d2 |
|  0  0 a3 b3 |   | x3 |   | d3 |
  a — поддиагональ, b — диагональ, c — наддиагональ

Как устроена прогонка

Прямой проход («прямая прогонка») последовательно исключает поддиагональ, выражая каждое неизвестное через следующее: $x_i = \alpha_i - \beta_i \cdot x_{i+1}$ (хранятся прогоночные коэффициенты). Обратный проход («обратная прогонка») идёт с конца, подставляя уже найденные значения. Всего два прохода по массиву — никаких вложенных циклов.

def прогонка(a, b, c, d):
    # a: поддиагональ (a[0] не используется), b: диагональ,
    # c: наддиагональ (c[-1] не используется), d: правая часть
    n = len(d)
    cp = [0.0] * n   # модифицированная наддиагональ
    dp = [0.0] * n   # модифицированная правая часть
    cp[0] = c[0] / b[0]
    dp[0] = d[0] / b[0]
    for i in range(1, n):           # прямой проход
        m = b[i] - a[i] * cp[i - 1]
        cp[i] = c[i] / m if i < n - 1 else 0.0
        dp[i] = (d[i] - a[i] * dp[i - 1]) / m
    x = [0.0] * n
    x[-1] = dp[-1]
    for i in range(n - 2, -1, -1):  # обратный проход
        x[i] = dp[i] - cp[i] * x[i + 1]
    return x

# Система 4x4, трёхдиагональная
a = [0, 1, 1, 1]    # поддиагональ
b = [2, 2, 2, 2]    # диагональ
c = [1, 1, 1, 0]    # наддиагональ
d = [1, 2, 3, 4]    # правая часть
x = прогонка(a, b, c, d)
print("решение:", [round(v, 6) for v in x])

Вывод:

решение: [0.0, 1.0, 0.0, 2.0]

Когда прогонка устойчива

Прогонка без перестановок устойчива, если матрица обладает диагональным преобладанием: $|b_i| \ge |a_i| + |c_i|$ для всех строк. К счастью, матрицы из сплайнов и краевых задач обычно ему удовлетворяют. Если преобладания нет, прямой проход может делить на близкое к нулю $m$, и нужна более общая (и более дорогая) схема с выбором главного элемента.

Как работает под капотом

Прогонка — это в точности LU-разложение трёхдиагональной матрицы, но записанное компактно: $L$ и $U$ сами получаются двухдиагональными, поэтому хранить и обрабатывать их можно одномерными массивами. Память тоже линейна — $O(n)$ вместо $O(n^2)$ у плотной матрицы. Именно это сочетание (линейное время и память) делает прогонку рабочим инструментом для систем с миллионами неизвестных в одномерных задачах. В библиотеке — scipy.linalg.solve_banded.

Частые ошибки

Хранить трёхдиагональную матрицу как плотную $n \times n$. Это убивает выигрыш по памяти; держите три массива $a, b, c$.
Путать индексацию диагоналей. $a[0]$ и $c[n-1]$ не существуют (выходят за матрицу) — следите за границами, иначе сдвиг на 1 даст неверный ответ.
Применять прогонку без диагонального преобладания. Можно получить деление на $\approx 0$; проверяйте условие устойчивости или берите общий метод.

Итоги

Прогонка (Томас) решает трёхдиагональную систему за $O(n)$ времени и памяти — двумя проходами.
Это компактный LU для матрицы из трёх диагоналей; общий Гаусс здесь избыточен.
Устойчива при диагональном преобладании $|b_i| \ge |a_i| + |c_i|$.
Незаменима для сплайнов, краевых задач и неявных разностных схем.