[백준 30413] 양 한 마리... 양 A마리... 양 A제곱마리... - python
원본 문제 링크 : 양 한 마리… 양 A마리… 양 A제곱마리…
문제
춘배는 오늘도 열심히 활동을 마친 후 자려 하지만 도저히 잠이 안 온다. 춘배는 잠을 자기 위해 자신만의 방법으로 양을 세려는데 한 마리씩 세지 않고 여러 마리를 한꺼번에 세면서 자려 한다.
춘배는 일단 양의 정수 $A$를 정한다. 그 후 양을 셀 때 첫 번째에 센 양의 수는 항상 1$1$로 두고, 그 뒤 두 번째에 센 양의 수는 $A$, 세 번째에 센 양의 수는 $A^{2}$ 이렇게 점점 양의 수를 세어 간다. 즉, $n$번째에 센 양의 수는 $A^{n-1}$ 가 된다.
춘배는 이러한 방식으로 양을 세다 문득 자신이 첫 번째부터 $B$번째까지 센 모든 양의 수가 얼마나 될지 궁금해졌다. 춘배를 위해 첫 번째부터 마지막 $B$번째까지 센 모든 양의 수가 몇 마리인지 구해보자! 하지만 수가 너무 커질 수 있기에 1\,000\,000\,007(= 10^{9} + 7)$로 나눈 나머지를 구하자.
입력
첫 번째 줄에 양의 정수 $A$, $B$ 가 공백으로 구분되어 주어진다. $(1 \le A \le 1\,000$, $ 1 \le B \le 10^{12})$
출력
첫 번째부터 $B$번째까지 센 모든 양의 수를 $1\,000\,000\,007$로 나눈 나머지를 출력한다.
힌트 : $1\,000\,000\,007$은 소수다.
풀이
총 양의 수?
만약 B번째까지 양을 샌다면, 총 양의 수를 수식으로 어떻게 나타낼 수 있을까?
\[1+A+A^2+A^3+\cdots+A^{B-1}\\=A^0+A+A^2+A^3+\cdots+A^{B-1}\]만약 등비수열에 대해 알고 있다면, 위의 식을 등비수열의 합 공식을 이용해 나타낼 수 있다.
\[A^0+A+A^2+A^3+\cdots+A^{B-1}=\frac{A^0(A^B-1)}{A-1}\]이때 A=1인 경우에 대해 예외처리를 해줘야하는 것을 잊지 말자.
수가 너무 커지는 것을 방지하기 위해서 $1000000007$ 로 나눈 나머지를 계산해야하는데, 최종 결과에만 나머지 연산을 적용하는 경우, 언어에 따라 연산이 느려지는 문제가 발생할 수 있다.
따라서 모듈러 연산 문제는 모듈러 연산의 분배법칙을 통해 연산 도중에 모듈러 연산을 적용하여 오버플로우가 나지 않도록 유지하는 것이 핵심이다.
Int 등의 자료형의 범위는 제한적이기에, 오버플로우를 방지하기 위해 모듈러 연산을 사용한다. 파이썬이 아닌 다른 언어는 정수에 큰 제약이 있다. 파이썬도 기본적으로 수가
21억*21억*2보다 커지면 계산과정이 느려진다.왜 $1000000007$를 모듈러 연산에 사용하는가?
- 소수이다. 즉 FlT를 사용하여 나눗셈에 대한 모듈러 연산의 분배법칙을 사용가능
- 모듈러가 지나치게 작다면 언어의 표현력을 비효율적으로 사용하는 것이므로, int의 표현 범위를 최대한 활용하기 위해 int의 max값에 가까운 소수를 사용하는 것
모듈러 연산에 관한 자세한 내용은 구현 코드 이해에도 중요하니, 이전에 정리한 포스트를 참고하는 것이 좋을 것 같다.
시간 제한에 비해 너무 많은 연산량
우리가 계산해야하는 최종 식에 존재하는 $A^B$는 이 문제의 시간복잡도를 증가시키는 가장 큰 요인이다.
따라서 해당 부분에서의 시간 복잡도를 줄이면, 시간 제한 내로 테스트케이스를 통과할 수 있을 것으로 보인다.
1
2
3
4
5
6
7
8
9
def mul(Num,n) : #분할 정복을 이용한
if n == 1 :
return Num
else :
x = mul(Num,n//2)
if n%2 == 0 :
return (x*x)%M
else :
return (x*x*Num)%M
위의 소스코드가 이해가 안된다면, 이전에 분할 정복 기반 거듭 제곱 최적화에 대해 정리한 포스트가 있으니 참고하는 것이 좋겠다.
소스 코드
최종적으로 구현하여 통과한 코드는 다음과 같다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
A, B = map(int,input().split())
M = 1000000007
def mul(Num,n) : #분할 정복을 이용한
if n == 1 :
return Num
else :
x = mul(Num,n//2)
if n%2 == 0 :
return (x*x)%M
else :
return (x*x*Num)%M
if A == 1 :
print(B%M)
else :
print((((mul(A,B)-1)%M)*(mul(A-1,M-2)%M))%M) #페르마의 소정리를 이용한 모듈러 연산