(Atcoder)ABC179 E - Sequence Sumをpythonで

備忘録です。

atcoder.jp

数え上げイメージ

コード

n,x,m = map(int,input().split())
id = [-1]*m
length = 0
a = []
tot = 0 # total
while id[x] == -1:
    a.append(x)
    id[x] = length # idにはxが現れたときの配列長さを格納
    length += 1
    tot += x
    x = (x*x)%m

# c:繰り返し配列長さ
# 繰り返しが確認できた配列長さ - 繰り返し配列開始位置    
c = length - id[x] 
# s 繰り返し配列の総和
s = 0 
for i in range(id[x],length):
    s += a[i]
ans = 0
if n <= length:
    for i in range(n):
        ans += a[i]
else:
    ans += tot
    n -= length
    ans += s*(n//c)
    n %= c
    for i in range(n):
        ans += a[id[x]+i]
print(ans)