Branched Evolution

Competitive Programming in Python

ABC 179 E - Sequence Sum

$a _ 1 = x, \ a _ {n + 1} = a _ {n} ^ 2\ (\mathrm{mod}\ m)$ で定まる数列 ${ a _ n }$ について,第 $n$ 項までの和を求める.($n \leq 10 ^ {10}, m \leq 10 ^ {5}$)

atcoder.jp

解説

数列は第 $0$ 項から始まることにする.$a _ {0} = x$,$S _ {n} = \sum _ {i = 0} ^ {n - 1} a _ {i}$ とする. $a _ {n}$ は $0$ から $m - 1$ までの $ m $ 通りの値しかとらないから,$0 \leq i < j \leq m $ かつ $a _ {i} = a _ {j}$ となるような $i,j$ が存在する. $n - i$ を $j - i$ で割った商と余りを $q, r$ とすると, $$ \begin{aligned} S _ {n} &= S _ {i} + (S _ {j} - S _ {i}) \times q + S _ {i + r} - S _ {i} \\ &= (S _ {j} - S _ {i}) \times q + S _ {i + r} \end{aligned} $$

実装

Submission #16925511 - AtCoder Beginner Contest 179
xを更新しながら配列Fxが現れたインデックスを記録し,ループを検出する.

n, x, m = map(int, input().split())
S = [0]
F = [None] * m
for j in range(m + 1):
    S.append(S[-1] + x)
    if F[x] is not None:
        i = F[x]
        q, r = (n - i) // (j - i), (n - i) % (j - i)
        print((S[j] - S[i]) * q + S[i + r])
        exit()
    F[x] = j
    x = pow(x, 2, m)