이항계수를 빠르게 구하는 알고리즘

안녕하세요? ryul입니다.

오늘은 이항계수를 빠르게 구하는 알고리즘에 대해서 공부해보겠습니다.

알고나면 매우 쉽습니다.

먼저 이항계수란?

고등학교 때 배웠던 조합이랑 같은 거라고 생각하시면 됩니다. 

N개의 공 에서 K개의 공을 뽑는 경우의수 를 의미하죠~

기호로는 어떻게 표시할까요?

nCk 

이렇게~

자! 먼저 가장 나이브한 방법으로 nCk 를 구한다면?

Phase 1 나이브

nCk = n! / (k!(n-k)!) 을 직접 계산으로 구하게 됩니다.

이 때 n이 커지게되면 nCk의 값은 분명 기하급수적으로 커지게 되므로 n이 클 경우 분명 문제에서 값을 p로 나눈 나머지를 구하라는 요구를 할 것입니다.

따라서 nCk mod p 의 값을 구해야 하는 것인데, 

단순히 (분모 mod p) / (분자 mod p) 의 값은 우리가 구하고자 하는 값이 아닙니다.

왜냐하면 우리는 나누어 떨어지는 경우 분자와 분모를 약분해야 하는데 단순히 n! mod p 를 구한 후에 하는 계산은 의미가 없습니다.

따라서 일일히 나누어 떨어지는 것부터 계산하게 될 경우 시간이 많이 걸리게 됩니다.

Phase 2 DP

n+1Ck+1 = nCk + nCk+1 이라는 식이 성립하는 것을 알 수 있습니다.

(직접 계산해보면 납득하실 거에요!)

따라서 메모이제이션을 이용하여 dp[n][k] 의 배열을 잡고 동적으로 계산해나간다면 O(n^2)의 시간복잡도로 이항계수를 구할 수 있습니다!

하지만 n이 너무커서 dp로도 되지 않는다면??

phase 3 페르마의 소정리

페르마의 소정리를 이용하겠습니다.

페르마는 말했습니다.

" p를 소수라고할 때 a^(p-1) = 1 (mod p) 이다."

따라서 a * a^(p-2) = 1 (mod p) 라는 것을 알 수 있기 떄문에 우리는 a의 역원이 a^(p-2)라는 것을 알 수 있게됩니다!

즉, 1/a mod p == a^(p-2) 가 되는 것이죠.

따라서 우리가 구하고자 했던 nCk = n! /(k! * (n-k)!) 에서 1 / (k! * (n-k)!) mod p 은 (k! * (n-k)!)^(p-2) mod p 라는 것을 알 수 있죠!!

자! 그렇다면 nCk mod p = ( n! mod p) * ( (k! * (n-k)!)^(p-2) mod p )  입니다.

식이 너무 더러우신가요? 자세히보면 충분히 이해할 수 있습니다.

Phase 4 전처리 후 O(1) 에 사용

n * inverse(n!) = inverse((n-1)!) 라는 식이 성립합니다. 

( 왜 성립하는 지는 식을 계산해보시면 알 수 있습니다!)

따라서 n! 의 inverse 만 분할정복을 통하여 O(log p) 의 시간에 한 번 구해준다면?

메모이제이션을 통하여 n!의 inverse 를 O(n) 의 구할 수 있습니다!

따라서 O(n + logp) 시간 복잡도로 이항 계수 전처리가 수행됩니다.

그 후에는 O(1)의 시간복잡도로 모든 nCk 를 구할 수 있게됩니다.

Phase4 까지 완료한 코드는 다음과 같습니다.

#include <cstdio>
#include <algorithm>
#define P 1000000007LL
typedef long long ll;
using namespace std;
ll fac[4000001], n, k, inv[4000001];//inv[x]의 정의는 x의inverse가 아닌 x!의 inverse
ll power(ll x, ll y) {    //분할 정복을 이용하여 x^y 구하기
    ll ret = 1;
    while (y > 0) {
        if (y % 2) {
            ret *= x;
            ret %= P;
        }
        x *= x;
        x %= P;
        y /= 2;
    }
    return ret;
}
int main() {
    fac[1] = 1;
    for (int i = 1; i <= 4000000; i++)
        fac[i] = (fac[i - 1] * i) % P;            //factorial 전처리
    inv[4000000] = power(fac[4000000], P - 2);    //페르마의 소정리를 이용하여 fac(400만)의 inverse를 구함(이때 분할 정복을 이용하여 logP의 시간에 처리) 
    for (int i = 4000000 - 1; i >= 0; i--) 
        inv[i] = (inv[i + 1] * (i + 1)) % P;    //inverse(fac(i))=inverse(fac(i+1))*(i+1)이라는 수식을 이용하여 전처리
    //총 O(N+logP)시간에 전처리를 끝냄 
    //전처리가 끝났기 때문에 어떤 쿼리가 들어와도 상수시간에 이항 계수를 계산 가능
    scanf("%lld%lld", &n, &k);
    if (n == k || !k) {
        puts("1");
        return 0;
    }
    ll r = (fac[n] * inv[n - k]) % P;
    r = (r*inv[k]) % P;
    printf("%lld\n", r);
    return 0;
}
출처:[ACM-ICPC 상 탈 사람]

이상 이항계수를 빨리 구하는 법에 대해서 알아봤습니다.

이제 경우의수 식을 잘 세우는 테크닉만 키우면 이 코드를 잘 이용해서 문제를 해결할 수 있습니다!

이상~