CP Notebook

← all snippets

Tridiagonal

x=tridiagonal(d,p,q,b) solves the equation system \[ ( cb_0 b_1 b_2 b_3 ⋮ b_n-1 ) = ( cccccc d_0 & p_0 & 0 & 0 & … & 0 q_0 & d_1 & p_1 & 0 & … & 0 0 & q_1 & d_2 & p_2 & … & 0 ⋮ & ⋮ & ddots & ddots & ddots & ⋮ 0 & 0 & … & q_n-3 & d_n-2 & p_n-2 0 & 0 & … & 0 & q_n-2 & d_n-1 ) ( cx_0 x_1 x_2 x_3 ⋮ x_n-1 ). \] This is useful for solving problems on the type \[ a_i=b_ia_i-1+c_ia_i+1+d_i, 1≤ i≤ n, \] where a_0, a_n+1, b_i, c_i and d_i are known. a can then be obtained from \a_i\=tridiagonal(&\1,-1,-1,...,-1,1\, \0,c_1,c_2,…,c_n\, &\b_1,b_2,…,b_n,0\, \a_0,d_1,d_2,…,d_n,a_n+1\). Fails if the solution is not unique. If |d_i| > |p_i| + |q_i-1| for all i, or |d_i| > |p_i-1| + |q_i|, or the matrix is positive definite, the algorithm is numerically stable and neither tr nor the check for diag[i] == 0 is needed.

Time: O(N) 26 lines Brute-force tested mod 5 and 7 and stress-tested for real matrices obeying the criteria above.

content/numerical/Tridiagonal.h — Ulf Lundstrom, Simon Lindholm, source: https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm

typedef double T;
vector<T> tridiagonal(vector<T> diag, const vector<T>& super,
		const vector<T>& sub, vector<T> b) {
	int n = sz(b); vi tr(n);
	rep(i,0,n-1) {
		if (abs(diag[i]) < 1e-9 * abs(super[i])) { // diag[i] == 0
			b[i+1] -= b[i] * diag[i+1] / super[i];
			if (i+2 < n) b[i+2] -= b[i] * sub[i+1] / super[i];
			diag[i+1] = sub[i]; tr[++i] = 1;
		} else {
			diag[i+1] -= super[i]*sub[i]/diag[i];
			b[i+1] -= b[i]*sub[i]/diag[i];
		}
	}
	for (int i = n; i--;) {
		if (tr[i]) {
			swap(b[i], b[i-1]);
			diag[i-1] = diag[i];
			b[i] /= super[i-1];
		} else {
			b[i] /= diag[i];
			if (i) b[i-1] -= b[i]*super[i-1];
		}
	}
	return b;
}