[백준 11401] 이항 계수 3

 

 

 수학에 분할 정복이 더해진 문제다. 이항 계수 문제 자체는 PS에서 자주 등장하는 문제이지만, 보통은 파스칼 항등식 개념을 메모이제이션과 분할 정복으로 구현해서 풀 수 있느냐를 묻는다. 하지만 이 문제는 N과 K의 범위가 커서 메모이제이션을 활용하면 메모리 초과로 풀 수 없게 된다. 그렇다고 메모이제이션 없이 분할 정복만 사용하는 것은 시간 복잡도가 끔찍해지고, 조합론을 활용해서 풀려하면 분모든 분자든 오버플로우가 날 수 있다(모듈러 연산을 분자와 분모에 분배하는 것은 등식이 성립하지 않는다).

 

 그래서 처음 시도했던 것은, 메모이제이션을 2차원 배열로 사용하지 말고 2차원 맵(unordered_map)을 사용해서 불필요한 공간 사용을 줄이는 것이었다. 거기에 이항 계수의 성질을 이용해서 k가 n/2보다 크면 k대신 n-k를 사용해 좀 더 최적화를 시도해보았으나 결국 메모리 초과를 받았다.

 

 이 문제는 한참 감을 못 잡다가, 결국 밑에 적힌 알고리즘 분류를 참고하게 되었다. 페르마의 소정리는 하루를 고민해도 떠올리지 못했을 것 같아서 보길 잘했다는 생각이 들었다.

 일단 위에서 설명했던 이유대로, 파스칼 항등식을 이용한 분할 정복으로는 풀 수 없기 때문에 조합론에서 일반적으로 이항계수를 구하는 식을 사용해야한다.

 결국 저 식을 활용해야 하는데, n과 k의 범위가 너무 커서 오버플로우가 발생할 수 있다. 분모와 분자에 각각 모듈러 연산을 한 값은 실제 값에 모듈러 연산의 취한 값과는 달라서 저 꼴을 변형할 필요가 있었다. 일단 분자의 식을 A로 두고 분모의 식을 B로 둔 상태로 구하고자 하는 것은 다음과 같다.

 

 

 이제 이 때 페르마의 소정리를 적용할 수 있다. 페르마의 소정리와 합동식의 성질에 의해 다음의 식들이 성립한다.

 

p가 소수이며, a가 p의 배수가 아닐 때.

 

 

 따라서 A와 B에 관한 식을 이렇게 바꾼다.

 

 식 B의 지수가 음수가 아니게 되었으므로, 위와 같이 이제 모듈러를 분배해서 오버플로우를 방지할 수 있다. 식 A와 B 자체는 그냥 반복문을 통해서 O(n)으로 쉽게 구해낼 수 있으나, 식 B에 대한 거듭제곱이 문제다. P 값이 매우 크기 때문인데, 이걸 단순하게 B를 반복적으로 곱하고 있으면 당연히 시간 초과가 난다.

 이 때 활용할 수 있는게 분할 정복이다. 이건 거듭제곱을 구할 때 자주 쓰이는 유형이기도 하다. 지수의 성질에 의해 a^(n+m)은 a^n * a^m이니까, a^2n은 a^n * a^n이다. 이를 이용해 지수가 절반인 거듭제곱을 구하는 문제로 계속 분할하고, 지수가 절반인 거듭제곱을 구했다면 그 값을 두번 곱해주면 된다.

 종료 조건을 지수가 0일 때 1을 리턴하도록 하고, 지수가 홀수일 때는 지수/2인 거듭제곱을 구해 두 번 곱하고 밑을 한 번 더 곱해서 본래의 지수를 맞춰주도록 하면 편하게 구현할 수 있다.

 

 이제 거듭제곱 문제도 처리 되었으므로, 위 과정들을 코드로 구현하면 오버플로우도 방지하면서 시간 초과 및 메모리 초과 없이 이 문제를 해결할 수 있게 된다.

#include <iostream>
#define MOD 1000000007

using namespace std;

long long power(int n, int exp) {
    if(exp == 0) return 1;

    long long half = power(n, exp/2) % MOD;
    if(exp%2==0) {
        return half * half % MOD;
    } else {
        return half * half % MOD * n % MOD;
    }
}

int binomialCoefficient(int n, int k) {
    if(k==0 || k==n) return 1;
    if(k==1) return n;

    long long a = 1, b = 1;
    
    for(int i=n; i>=n-k+1; i--) {
        a = (a*i)%MOD;
    }
    for(int i=1; i<=k; i++) {
        b = b*i%MOD;
    } 
    
    return a * power(b, MOD-2) % MOD;
}

int main() {
    int n, k;

    cin >> n >> k;

    cout << binomialCoefficient(n, k);
}