Modulo Prime Power Trick

이전에 조합을 소수의 $k$제곱으로 나눈 나머지를 구하는 방법을 소개했다. 그 방법은 너무 어렵기 때문에 조금 더 쉬운 방법을 소개한다.

 

- Notation

$(u!)_{p}$는 $u!$에서 $p$의 배수인 것을 제외한 곱을 의미한다.

 

소수 $p$와 자연수 $k$에 대하여 $p^{2k}$으로 나눈 나머지를 구한다고 하자. 이때 다음 식을 이용한다.

$$\frac{(dp^{k}!)_{p}}{((d-1)p^{k}!)_{p}} \equiv (p^{k}!)_{p} \pmod{p^{2k}}$$

풀어서 써보면 다음과 같다. $p^{k}$와 서로소인 $p^{k}$이하의 자연수들의 집합을 $\{S\}$라고 하자. 

$$\prod_{a \in S} (dp^{k}+a) \equiv \prod_{a \in S} a \pmod{p^{2k}}$$

이것을 증명해보자! 좌변의 식을 풀면 다음과 같이 쓸 수 있다.

$$p^{2k}(...)+dp^{k}\left(\prod_{a \in S} a\right)\left(\sum_{a \in S} \frac{1}{a}\right)+\prod_{a \in S}a$$

$\prod_{a \in S} a$는 전부 $p$와 서로소인 수들의 곱이므로 $p$로 나누어 떨어질 수 없다. 따라서 합 부분이 $p^{k}$으로 나누어 떨어짐을 보이면 된다. $p$가 2보다 큰 경우와 2인 경우를 나누어서 살펴본다.

 

• $p>2$

$1<a,\ b \in S<p^{k}-1$인 자연수 $a,b$에 대해 $ab \equiv 1 \pmod{p^{k}}$ 이 성립한다고 하자. 윌슨의 정리의 확장에 의해 이런 방식으로 집합 $S$에 있는 원소 $x$($1<x \in S<p^{k}-1$)를 묶을 수 있다. 1과 $p^{k}-1$은 약간 다르게 표현된다. 이는 추후 설명한다.

$$\sum_{a \in S} \frac{1}{a}$$

에서 앞서 설명한 $a,b$를 선택해서 더해보자.

$$\frac{1}{a}+\frac{1}{b}=\frac{a+b}{ab} \equiv -(a+b) \pmod{p^{k}}$$

$x=1, p^{k}-1$일 때 $x^{2} \equiv 1 \pmod{p^{k}}$이므로 $\frac{1}{x} \equiv x \pmod{p^{k}}$이다. 

따라서 다음이 성립한다.

$$\sum_{a \in S} \frac{1}{a} \equiv -\sum_{a \in S} a \pmod{p^{k}}$$

$$=-\frac{(1+p^{k})p^{k}}{2}+\frac{(p+p^{k})p^{k-1}}{2}=-\frac{p^{2k-1}(p-1)}{2}$$

$p-1$이 2로 나누어 떨어지므로 

$$\sum_{a \in S} \frac{1}{a} \equiv \sum_{a \in S} a \equiv 0 \pmod{p^{k}}$$

 

• $p=2$

$$\sum_{a \in S} \frac{1}{a} \equiv -\sum_{a \in S} a \pmod{2^{k}}$$

$$=-\frac{(1+2^{k})2^{k}}{2}+\frac{(2+2^{k})2^{k-1}}{2}=-2^{2k-2}$$

에서 $2k-2 \ge k$,  $k \ge 2$를 얻는다. 따라서

$$\sum_{a \in S} \frac{1}{a} \equiv -\sum_{a \in S} a \equiv 0 \pmod{2^{k}}$$

$p=2$인 경우 $k \ge 2$임을 주의해서 코드를 짜자.

 

따라서 다음이 성립한다.

$$\prod_{a \in S} (dp^{k}+a) \equiv \prod_{a \in S} a \pmod{p^{2k}}$$

 

이제 이 식을 이용해서 $(N!)_{p} \pmod{p^{2k}}$를 계산할 것이다. 알고리즘은 다음과 같다.

1. $p^{k}$이하의 $p^{k}$과 서로소인 수들의 곱$s$과 $s^{2},\ ...,\ s^{p^{k}}$을 전처리한다. 

2. $N$을 $p^{2k}$로 나눈 나머지를 $t$라고 하자. 답에 $s^{t/p^{k}}$를 곱한다. -1을 아래에 있는 설명을 참고해 적절히 곱한다.

3. $N$을 $p$로 나눈 후 2번 과정을 $N$이 0이 될 때까지 반복한다.

이 알고리즘의 작동 원리를 설명하겠다.

$N$을 $p^{2k}$로 나누는 이유를 알아보자. $p^{k}$과 $p^{k}$이하이고 $p^{k}$과 서로소인 수들의 곱은 $\pm 1$과 합동임을 윌슨의 정리의 확장을 통해 알고 있다. $\pm$은 $p=2$일 때 $+$이고 나머지는 $-$이다. $N$이라는 수를 $p^{2k}$ 단위로 나누어 생각해보면 아래와 같다.

$$(N!)_{p}=[1\times...\times(p^{2k}-1)][(p^{2k}+1)\times...\times(2p^{2k}-1)]...[(up^{2k}+1)\times...\times(up^{2k}+t)]$$

$$\equiv (\pm 1)^{u}[1\times...\times t]\pmod{p^{2k}}$$

1부터 $t$까지의 수들 중 $p^{k}$으로 묶이는 수들의 곱은 전처리 되어있는 상태이다. 또한 앞서 증명한 식에 의해 다음과 같이 쓸 수 있다.

$$\equiv (\pm 1)^{u}[1\times...\times (p^{k}-1)][(p^{k}+1)\times...\times (2p^{k}-1)]...[(v p^{k}+1)\times...\times (vp^{k}+r)]\pmod{p^{2k}}$$

$$\equiv (\pm 1)^{u}s^{v}[(vp^{k}+1)\times...\times (vp^{k}+r)]\pmod{p^{2k}}$$

$[(vp^{k}+1)\times...\times (vp^{k}+r)]$이 부분은 따로 계산을 해주면 된다.

 

조합 식을 변형하자. $v_{p}(n)$을 $p^{k} |n$을 만족하는 최대의 $k$로 정의하자.

  $${n \choose m}=p^{v_{p}(n)-v_{p}(m)-v_{p}(n-m)} \frac{(n!)_{p}}{(m!)_{p}((n-m)!)_{p}}$$

분모에 있는 수들은 어차피 $p$와 서로소이기 때문에 모듈러 역원을 구해 곱해주면 된다.

 

이것을 구현한 코드이다. 

코드 보기

#include<stdio.h>
#include<algorithm>
#include<utility>
#define MAX 1000001
using namespace std;
typedef long long ll;
typedef pair<ll, ll> Pair;
ll fpow(ll n, ll k, ll mod = 0){
    ll s = 1;
    for(;k;k >>= 1){
        if(mod){
            if(k & 1) s = (__int128) s * n % mod;
            n = (__int128) n * n % mod;
        }
        else{
            if(k & 1) s = s * n;
            n = n * n;
        }
    }
    return s;
}
ll power[MAX] = {1};
Pair factorial_p(ll n, ll p, ll pq, ll bpq){
    ll ret = 1, pt = 0, t, i;
    while(n){
        t = n % pq;
        if(p != 2 && (n / pq) % 2) ret = pq - ret;
        ret = (__int128) ret * power[t / bpq] % pq;
        for(i = t / bpq * bpq + 1;i <= t;i++){
            if(i % p == 0) continue;
            ret = (__int128) ret * i % pq;
        }
        pt += n / p;
        n /= p;
    }
    return make_pair(ret, pt);
}
ll bincoeff_pq(ll n, ll m, int p, int q){
    ll real = fpow(p, q);
    if(q % 2) q += 1; //match with even number
    ll pq = fpow(p, q), bpq = fpow(p, q / 2);
    if(p == 2){
        pq = fpow(p, max(4, q));
        bpq = fpow(p, max(4, q) / 2);
    }
    ll gop = 1, i;
    Pair get;
    ll up, dn, pt, ans;
    for(i = 1;i < bpq;i++){
        if(i % p == 0) continue;
        gop = (__int128) gop * i % pq; 
    }
    for(i = 1;i <= bpq;i++){
        power[i] = (__int128) power[i - 1] * gop % pq;
    }
    get = factorial_p(n, p, pq, bpq);
    up = get.first; pt = get.second;
    get = factorial_p(m, p, pq, bpq);
    dn = get.first; pt -= get.second;
    get = factorial_p(n - m, p, pq, bpq);
    dn = (__int128) dn * get.first % pq; pt -= get.second;
    if(pt >= q) return 0;
    ans = (__int128) up * fpow(dn, pq / p * (p - 1) - 1, pq) % pq;
    return (__int128) ans * fpow(p, pt) % real;
}
 
int main(){
    ll n, m;
    int p, q;
    scanf("%lld %lld %d %d", &n, &m, &p, &q);
    printf("%lld", bincoeff_pq(n, m, p, q));
}