[백준] 이항 계수 3(페르마 소정리, modular inverse, 분할정복)

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

11401번: 이항 계수 3

페르마 소정리를 이용한 모듈러 역원과 분할정복을 활용하는 문제입니다. 

 

문제부터 간단히 요약해보면, 

nCk의 값을 1,000,000,007로 나눈 나머지를 구하라는 문제입니다. 

 

굉장히 간단해 보이지만, nCk의 값에서 n의 범위가 4,000,000 까지 가능하므로 그냥 직접 계산을 할 시 바로 시간초과가 되게 됩니다. 아이디어와 주의해야하는 점에 대한 이야기를 해 보도록 하겠습니다. 

 

<Solution>

우선 아이디어는 다음과 같습니다. 

다음을 활용하여 식을 정리한 후에 그것을 활용하여야 합니다. 

다음과 같이 식을 정리해줍니다. 그러면 우리는 nCk와 n!(k!(n-k)!)^(p-2)의 mod p 연산 결과가 같음을 알 수 있습니다. 

이제 이것을 활용하여 문제를 해결 할 수 있습니다. 

 

이 때 주의해야하는 점은 factorial 연산을 할 때, n이 4,000,000을 최댓값으로 가지기에 n!를 통으로 구하면 안됩니다. 

factorial 이라는 함수를 만들고 매번 %p를 함으로써, 결과를 항상 일정범위 안으로 유지 할 수 있습니다. 이렇게 해야 timeout 되지 않고 백준의 testcase들을 모두 통과할 수 있습니다. 

또한 (k!(n-k)!)^(p-2) 를 구할때에는 분할정복 방식을 사용하여 시간을 단축하여야 합니다.

 

실제 구현은 아래와 같습니다.

import sys
import math

p = int(1e9)+7
n,k = map(int,(sys.stdin.readline().split()))

def factorial(n: int) -> int:
    result = 1
    for i in range(1,n+1):
        result *= i 
        result %=p
    return result


def exp(base: int, expo: int) -> int:
    if expo == 1:
        return base
    a = exp(base, expo//2)
    b = expo%2
    # print(f'a,b = {a,b}')
    return a*a*(base**b) %p

print((factorial(n))%p * exp((factorial(k))*(factorial(n-k)), p-2)%p)

만약 위의 코드의 factorial 함수에서 result %p 구문을 주석처리 하고 코드를 진행 하게 되면 항상 timeout이 됨을 알 수 있습니다.