[백준 BOJ_11401] 이항 계수 3 Python 풀이

출처: 백준 온라인 저지

문제

풀이

이항계수는 다음과 같은 공식으로 구해집니다.

\[\binom{N}{K}=\frac{N!}{K!(N-K)!}\]

이러한 이항계수 공식을 페르마의 소정리를 이용하여 다르게 쓸 수 있습니다.

페르마의 소정리란:

\[a^{p}\equiv a\ (mod\; p)\]

다시 말하자면 정수 $a$와 소수 $p$가 주어졌을 때, $a$의 $p$제곱을 $p$로 나눈 나머지는 $a$를 $p$로 나눈 나머지와 같다라는 정리입니다.

페르마의 소정리에서 양 변을 정수 $a$로 나눠줬을 때 다음과 같은 식을 얻을 수 있습니다.

$a^{p}\ (mod\; p) = a\ (mod\; p)$ (페르마의 소정리) \[a^{p-1}\ (mod\; p) = 1\ (mod\; p)\] \[a^{p-2}\ (mod\; p) = a^{-1}\ (mod\; p)\]

다시 본 문제로 돌아가 구해야 할 값을 이항계수 공식을 대입해 나타내보겠습니다. (편의상 $mod$를 $\%$로 대체합니다.)

\[\binom{N}{K}\ \%\ p=\frac{N!}{K!(N-K)!}\ \%\ p\] ($A=N!,\ B=K!(N-K)!$ 로 치환) \[\binom{N}{K}\ \%\ p=\frac{A}{B}\ \%\ p\] \[\binom{N}{K}\ \%\ p=AB^{-1}\ \%\ p\]

이 부분에서 모듈러 연산의 성질(Properties of modular arithmetic)을 적용하여 모듈러 연산을 분배해줍니다. 모듈러 연산의 성질이란:

1. $(a\ mod\ n+b\ mod\ n)\ mod\ n = (a+b)\ mod\ n$
2. $(a\ mod\ n-b\ mod\ n)\ mod\ n = (a-b)\ mod\ n$
3. $(a\ mod\ n*b\ mod\ n)\ mod\ n = (a*b)\ mod\ n$

3번 성질과 페르마의 소정리에서 얻은 마지막 식을 이용한다면 다음과 같이 나타낼 수 있다.

\[\binom{N}{K}\ \%\ p=\left(\left(A\ \%\ p\right)*\left(B^{-1}\ \%\ p\right)\right)\ \%\ p\] \[\binom{N}{K}\ \%\ p=\left(\left(A\ \%\ p\right)*\left(B^{\ p-2}\ \%\ p\right)\right)\ \%\ p\] \[\binom{N}{K}\ \%\ p=\left(\left(N!\ \%\ p\right)*\left(\left(K!\left(N-K\right)!\right)^{p-2}\ \%\ p\right)\right)\%\ p\]

factorial 함수의 경우 재귀함수로 구현 가능하지만, 입력 숫자가 너무 클 경우 Runtime error(Recursion error)가 났습니다. 같은 시간 복잡도라면 재귀를 쓰지 않고 for loop으로 구현해주는 것이 더 효율적으로 보입니다.

factorial 함수와 power 함수 모두 3번째 모듈러 연산의 성질을 이용해 매번 곱해질 때마다 소수 $p$의 값으로 모듈러 연산을 해주어 수가 커지지 않게 해줍니다. 모듈러 연산의 성질은 분배법칙과는 다르게 각각의 수의 모듈러 연산을 해준 뒤에 한번 더 해주어야 하기 때문에 함수 호출 후 잊지 않고 한번 더 모듈러 연산을 해줍니다.

코드

def factorial(n):
    ret = 1
    for i in range(1, n+1):
        ret = (ret*i)%p
    return ret
def power(n, exp):
    if exp == 0:
        return 1
    half = power(n, exp // 2)
    if exp % 2:
        return (half*half*n)%p
    return (half*half)%p
N, K = map(int, input().split())
p = 1000000007
print((factorial(N))*(power(factorial(K)*factorial(N-K), p-2))%p)