Official

E - Geometric Progression Editorial by en_translator


The sequence \((a_n)\) defined by \(a_n = \displaystyle\sum_{i=0}^{n-1} A^i\) satisfies \(a_0 = 0\) and \(a_{n+1} = Aa_n + 1\). What we want is \(a_X \bmod M\).

By the equation above, \(\begin{pmatrix} a_{n+1} \\ 1 \end{pmatrix} = \begin{pmatrix} A & 1 \\ 0 & 1 \end{pmatrix} \begin{pmatrix} a_n \\ 1 \end{pmatrix}\).

Since \(a_0 = 0\), we have \(\begin{pmatrix} a_X \\ 1 \end{pmatrix} = \begin{pmatrix} A & 1 \\ 0 & 1 \end{pmatrix}^X \begin{pmatrix} 0 \\ 1 \end{pmatrix}\).

The components of \(\begin{pmatrix} A & 1 \\ 0 & 1 \end{pmatrix}^X\), modulo \(M\), can be found in an \(O(\log X)\) time with the fast exponentiation (just as we find the \(X\)-th power of \(X\) modulo \(M\)).

Thus, the problem has been solved in a total of \(O(\log X)\) time.

Note that, if you use \(\displaystyle\sum_{i=0}^{X-1} A^i = \frac{A^X-1}{A-1}\), \((A-1)\) may not have an inverse in \(\text{mod } M\).

Sample code

#include <iostream>
#include <vector>
using namespace std;
using ll = long long;

vector<vector<ll>> mat_mul(vector<vector<ll>> a, vector<vector<ll>> b, ll mod) {
	// 行列乗算
	int n = a.size();
	vector<vector<ll>> res(n, vector<ll>(n));
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			for (int k = 0; k < n; k++) {
				res[i][j] += a[i][k] * b[k][j];
				res[i][j] %= mod;
			}
		}
	}
	return res;
}

vector<vector<ll>> mat_pow(vector<vector<ll>> a, ll b, ll mod) {
	// Matrix exponentiation
	int n = a.size();
	vector<vector<ll>> res(n, vector<ll>(n));
	for (int i = 0; i < n; i++) res[i][i] = 1;
	while (b) {
		if (b & 1) res = mat_mul(res, a, mod);
		a = mat_mul(a, a, mod);
		b >>= 1;
	}
	return res;
}

int main() {
	ll a, x, m;
	cin >> a >> x >> m;
	vector<vector<ll>> f = { {a,1},{0,1} };
	vector<vector<ll>> g = mat_pow(f, x, m);
	cout << g[0][1] << '\n';
}

posted:
last update: