[코딩 테스트 연습(파이썬/Python)] 백준(BOJ) 15712번 _ 등비수열

https://www.acmicpc.net/problem/15712

15712번: 등비수열

[문제]

첫 항이 a, 공비가 r, 항의 수가 n인 등비수열의 합을 mod로 나눈 나머지를 구하시오.

 

[입력]

첫째 줄에 a, r, n, mod가 공백으로 구분되어 주어진다. a, r, n, mod는 모두 1보다 크거나 같고, 109보다 작거나 같은 자연수이다. 

[출력]

문제의 정답을 출력하라.

 

[아이디어]

1) 등비수열의 합 공식을 활용하기는 어렵다.

2) 각 항의 합을 구하되, 항의 개수가 짝수일 때와 홀수일 때를 나누어 절반씩의 합을 분할 정복한다.

 

a, r, n, mod=map(int,input().split())
def f(r,n,mod):
    if n==1:
        return 1
    elif n%2==0:
        return (f(r,n//2,mod))*(pow(r,n//2,mod)+1)%mod
    else:
        return (f(r,n-1,mod)+pow(r,n-1,mod))%mod       
print(a*f(r,n,mod)%mod)

꽤 짧아서 조금 수정해 숏코딩으로도 만들어봤다.

def f(a,r,n,x):
    if n==1:return a%x
    elif n%2==0:return(f(a,r,n//2,x))*(pow(r,n//2,x)+1)%x
    else:return(f(a,r,n-1,x)+a*pow(r,n-1,x))%x       
print(f(*map(int,input().split())))