[BOJ 11401] 이항 계수 3 - Python 풀이

문제

 

입력

 

출력

 

코드

import sys
P = 1000000007

def power(a, b):
    if b == 0:
        return 1
    if b % 2:  
        return (power(a, b//2) ** 2 * a) % P
    else:
        return (power(a, b//2) ** 2) % P

def solve(N, K):
    fact = [1 for _ in range(N+1)]

    for i in range(2, N+1):
        fact[i] = fact[i-1] * i % P

    A = fact[N]
    B = (fact[N-K] * fact[K]) % P

    print((A % P) * (power(B, P-2) % P) % P )

if __name__ == "__main__":
    N, K = map(int, sys.stdin.readline().split())
    solve(N,K)

 

결과

 

 

간단한 문제이나 N의 크기가 40,000,000이나 되는게 함정인 문제였다.

n!/(n-k)k! 을 이용해서 풀려고 했으나, 1,000,000,007로 나눠야 하므로 분자와 분모에 나머지를 구해서 계산하게 되면 0으로 나누는 문제가 발생할 수 있기 때문에 분수꼴을 곱셈꼴로 바꿔야 한다.

(https://rhdtka21.tistory.com/123 <- 링크 참조)

 

따라서 페르마의 소정리를 이용해야 한다

페르마의 소정리는 p가 소수이고 a가 p의 배수가 아니면 a^p−1≡1 (모듈러 p) 라는 것이다.

즉 a^p-1을 p로 나눈 나머지는 1이라는 것.

(https://namu.wiki/w/%ED%8E%98%EB%A5%B4%EB%A7%88%EC%9D%98%20%EC%86%8C%EC%A0%95%EB%A6%AC <- 링크 참조)

 

이것을 이용하면

(n!/k!(n-k)!) % p = n!%p*(k!(n-k)!)^p-2%p 이고 이 식을 적용해야 한다. 

 

또한 제곱수를 구하는 방법은 분할정복을 이용하여 구해야 시간복잡도 O(logN)을 맞출 수 있다.