E - Geometric Progression Editorial by shakayami


答えを\(f(A,X,M)\)とします.

このとき、\(X=0\)ならば\(f(A,0,M)=0\) です。

さらに、\(X\)が奇数のときは

\(f(A,X,M)=1+Af(A,X-1,M)\)

または

\(f(A,X,M)=f(A,X-1,M)+A^{X-1} \mod M\)

となります。 さらに、\(X\)が偶数のときは(\(X=2n\))

\(f(A,2n,M)=f(A,n,M)\cdot (1+A^n) \mod M\)

または

\(f(A,2n,M)=f(A^2\mod M,n,M)\cdot (1+A) \mod M\) となります

以上をもとに再帰関数を書けばACするコードがかけます。

ここで、\(A^n \mod M\)は繰り返し二乗法を使うと\(O(\log{n})\)で求められることに注意です。

実装例1

def solve(A,X,M):
    if X==0:
        return 0
    if X%2==1:
        return (solve(A,X-1,M)*A+1)%M
    if X%2==0:
        return (solve(A*A%M,X//2,M)*(1+A))%M
a,x,m=map(int,input().split())
print(solve(a,x,m))

実装例2

def solve(A,X,M):
    if X==0:
        return 0
    if X%2==1:
        return (solve(A,X-1,M)+pow(A,X-1,M))%M
    if X%2==0:
        return (solve(A,X//2,M)*(1+pow(A,X//2,M)))%M
a,x,m=map(int,input().split())
print(solve(a,x,m))

posted:
last update: