【C++】漸化式と行列累乗

何の記事か

漸化式を行列に変換することで、より高速に高次の項を求めることができる「行列累乗」というアルゴリズムについて説明するとともに、2項間漸化式を例にとってC++での実装例を紹介する記事です。(この記事の内容は蟻本3-4節の内容に基づきます。)

2項間漸化式の行列累乗

例えば次のような問題を考えます。

• 数列${a_i}$が、$a_0=s,a_1=t$で、

a_{i+2} = pa_{i+1}+qa_i

のような関係式で記述されるとき、$n=10^{16}$に対して$a_n$を$998244353$で割ったあまりを求めてください。

こういった問題を、行列累乗を用いることで$O(logn)$で解くことができます。漸化式を行列を用いて表すと、

\begin{pmatrix}
a_{i+2} \\
a_{i+1}
\end{pmatrix}
=
\begin{pmatrix}
p &q \\
1 & 0
\end{pmatrix}
\begin{pmatrix}
a_{i+1} \\
a_{i}
\end{pmatrix}

となることが分かります。つまり、

\begin{pmatrix}
p &q \\
1 & 0
\end{pmatrix}
= A

とすると、

\begin{pmatrix}
a_{n+1} \\
a_{n}
\end{pmatrix}
=
A^n
\begin{pmatrix}
a_{1} \\
a_{0}
\end{pmatrix}

のようにして$a_n$を求めることができるのが分かります。ではどうやって$A_n$を求めるかというと、繰り返し二乗法を行列演算に拡張することで$O(logn)$で求まります。繰り返し二乗法は、$n$を$2$の累乗の和で表して、$i$乗の項を二乗することで$i+1$乗の項を求めることを言います。$n$をシフトするだけなので実装は非常に単純です。

const int M = 998244353;
typedef vector<int> vi;
typedef vector<vi> vvi;

//行列の乗算
vvi matrix_multiply(vvi X, vvi Y) {
	vvi Z(X.size(), vi(Y[0].size()));
	rep(i, X.size()) {
		rep(k, Y.size()) {
			rep(j, Y[0].size()) {
				Z[i][j] = (Z[i][j] + X[i][k] * Y[k][j]) % M;
			}
		}
	}
	return Z;
}

//A^nの計算
vvi matrix_pow(vvi A, ll n) {
	vvi B(A.size(), vi(A[0].size()));
	//単位行列でBを初期化
	rep(i, B.size()) {
		B[i][i] = 1;
	}

	while (n>0) {
		if (n & 1) { B = matrix_multiply(B, A); }
		A = matrix_multiply(A, A);
		n = n >> 1;
	}
	return B;
}

int main() {
	int n, s, t, p, q; cin >> n >> s >> t >> p >> q;
	vvi A(2, vi(2));
	A[0][0] = p; A[0][1] = q;
	A[1][0] = 1; A[1][1] = 0;
	A = matrix_pow(A, n);

	cout << (A[1][0] * t + A[1][1] * s) % M << endl;

	return 0;
}

m項間漸化式の行列累乗

一般のm項間漸化式

a_{i+m} = \sum_{j=0}^{m-1} p_j a_{i+j}

のような場合でも、

A=
\begin{pmatrix}
p_{m-1} & \cdots & p_1 & p_0 \\
1 & \cdots & 0 & 1 \\
\vdots & \ddots & \vdots & \vdots \\
0 & \cdots & 1 & 0
\end{pmatrix}

とすることで、同様の実装で動作します。行列の乗算部matrix_multiplyが重くなり、計算量は$O(m^3logn)$となります。