問題概要
整数 $N, X, M $ が与えられる．次の漸化式で定まる数列 ${ A_i }$ を考える．
\begin{align*}
A_n = \begin{cases}
X & \text{($n = 1$)} \\
( A_{ n - 1 } )^2 \bmod M & \text{(otherwise)}
\end{cases}
\end{align*}
$\sum_{ i = 1 }^N A_i$ を求めよ．
制約
- $1 \leq N \leq 10^{ 10 }$
- $0 \leq X < M \leq 10^5$
解法
数列を 1 項ずつナイーブに生成するコードは簡単に実装できますが，$\Theta( N )$ 時間となって TLE します．ちょっと考えると，$M $ の値が大きくないことから，数列に出現する値の種類も多くないことが分かります．もっと言えば，$M < N$ ならば，鳩ノ巣原理 的な理由で 2 回以上出現する値があります．漸化式は一意的に定まる形をしているので，2 回以上出現する要素があるのなら，その部分は周期性をもちます．
数列を 1 項ずつ生成して $N$ を $1$ ずつ減らしながら各値が何回出現するかをカウントしつつ周期の検出もするとします．ある値が $i$ 項目と $j$ 項目 ($i < j$) に出現したとすれば，周期の長さ $d$ は $d = j - i$ となります．また，周期が完全に含まれる回数 $t$ は，残りの $N$ を使って $t = \frac N d$ と定まります．そしてこの周期を一周舐めてあげると，周期に含まれる値の出現回数をまとめて計算することができます．その後で，残りの項数 $N$ を $N \bmod d$ で置き換えてあげれば，周期の処理が完了します．
$d$ で剰余をとることから，残りの項数は $d \leq M $ 以下となり*1，残りは愚直に処理できます．
この解法は，$M $ 通りある各値を，「周期に入る前」「周期の中」「周期を抜けた後」の高々 3 回処理すればよいので，$O( M )$ 時間となります．
コード
int main() { IN( LL, N, X, M ); VT< LL > counts( M ), last( M, -1 ), done( M ); LL a = X; for ( LL i = 0; N; ) { if ( last[a] == -1 || done[a] ) { ++counts[a]; last[a] = i; ( a *= a ) %= M; --N; ++i; continue; } const int d = i - last[a]; const LL t = N / d; LL aa = a; do { counts[ aa ] += t; ( aa *= aa ) %= M; done[ aa ] = true; } while ( aa != a ); N %= d; i += d * t; } LL res = 0; REP( i, M ) { res += LL( i ) * counts[i]; } cout << res << endl; return 0; }
