BJ 11041 이항 계수 3

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

페르마의 소정리

$p$가 소수이고, $a$가 $p$의 배수가 아닐 경우 다음이 성립합니다.

이를 페르마의 소정리라고 합니다. 이는 모든 소수가 만족하는 필요조건이지만, 충분조건은 아닙니다.

접근 방법

1부터 400만까지의 수라면 이전 DP에서의 $O(n^2)$의 방법으로는 턱없이 부족합니다. 적어도 $O(n\log n)$이나 $O(n)$의 수준까지 끌어올려야 합니다. 정수론을 이용해 다음과 같은 방법을 생각할 수 있을까요?

• 1부터 400만까지의 수에 대한 계승을 1,000,000,007로 나눈 나머지를 구합니다: $O(n)$
• $n$, $(n-k)$, $k$에 대한 계승값을 서로 나눕니다. 즉, $(n!\ \%\ m)/((n-k)!\ \%\ m)(k!\ \%\ m)$

아쉽게도 이는 두 번째 식이 성립하지 않기에 불가능합니다. 다만, 조금 다른 시각으로 바라볼 필요가 있습니다. 바로 역원을 구하는 것이지요.

$a\cdot a^{-1}\equiv 1 \pmod m$을 만족하는 $a^{-1}$을 법 $m$에 대한 $a$의 역원이라고 정의합니다.

1,000,000,007은 소수입니다. 페르마의 소정리에 의해 4,000,000 < 1,000,000,007이므로 1부터 400만까지의 $a$에 대해 $a^{1000000006}\equiv 1 \pmod{1000000007}$입니다. 따라서 양 변에 $a^{-1}$을 곱하면

입니다.

그렇다면 $k$와 $(n-k)$에 대해 각각 $k^{-1}\equiv k^{1000000005}$과 $(n-k)^{-1}\equiv(n-k)^{1000000005}$을 구하면 되겠네요.

빠른 거듭제곱의 계산

거듭제곱을 구하는 연산은 $O(\log n)$의 시간복잡도로 해결할 수 있습니다. Divide-and-conquer 방식을 이용하여

와 같이 2로 나눠가며 구할 수 있습니다.

본 계산에서는 이를 1,000,000,007로 나누어 가며 시행해야 합니다.

코드

def pow(n, k, m):
    if k == 1:
        return n
    pow_half = pow(n, k//2, m)
    if k % 2 == 0:
        return (pow_half ** 2) % m
    else:
        return (pow_half ** 2 * n) % m

n, k = map(int, input().split())

factorial = [1] * (n+1)
for idx in range(2, n+1):
    factorial[idx] = (factorial[idx-1] * idx) % 1000000007

def inverse(n):
    return pow(factorial[n], 1000000005, 1000000007)

print(factorial[n] * inverse(n-k) * inverse(k) % 1000000007)