백준 18151번: DivModulo

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

18151번: DivModulo

---

정수론에 대한 지식이 부족한 저로서는 매우 힘든 문제 중 하나였습니다. 처음에는 Lucas' Theorem과, Andrew Granville의 방법을 쓰려고 했지만, 이미 14854번: 이항 계수 6 및 19124번: Binomial Coefficient 에서 Granville의 방법을 구현해봤는데 그때 코드가 꽤 길었던 것이 기억에 남아 다른 블로그 글들을 참고하여 풀게 됐습니다. 결과적으로는 처음 시도하려던 게 더 구현이 편했을 수도 있겠다 싶습니다.

 

우선 저는 다음 세 글을 중점적으로 참고했음을 밝힙니다.

1. rkm0959, 12월 말 PS 정리 - Part 2
2. hapby9921, Modulo Prime Power Trick
3. Andrew Granville, 1997. Arithmetic Properties of Binomial Coefficients, I: Binomial Coefficients Modulo Prime Powers

---

일단 Hapby9921의 블로그 글을 먼저 리뷰하겠습니다.

 

글을 언뜻 보면, $(X!)_{p}$을 $p^{2k}$로 나눈 나머지를 구하는 방법만을 소개하고 있습니다. 그러나 홀수 지수에 대한 알고리즘을 고안할 필요는 없습니다.

 

$b | a$이면 $X \; \% \; b = (X \; \% \; a) \; \% \; b$입니다. 따라서 위 글에서는 일반적인 modulus $p^{q}$에 대해, $p^{2k} \geq p^{q}$인 $p^{2k}$을 잡아 $(X!)_{p} \; \textrm{mod} \; p^{2k}$을 먼저 구하겠다는 것입니다.

 

글의 핵심 중 하나인 $N = up^{2k}+t=up^{2k}+vp^{k}+r$로 놓고 $(N!)_{p}$를 전개하는 것은 쉽지만, 쉽게 구현하려면 $k$의 범위에 주목하는 게 편합니다.

 

Gauss가 Wilson's Theorem을 일반화했을 때, $(4!)_{2} \equiv -1 \; (\textrm{mod} \; 4)$ 입니다. 물론 이런 경우를 포함하여 구현하는 것이 전체적인 난이도를 크게 올리진 않겠습니다만, 다소 귀찮을 수는 있습니다. 그래서 위 글에서는 $p=2$일 때 특별히 $p^{2k} \geq 16$으로 놓고 있습니다. 이렇게 놓으면, $(2^{2k}!)_{p=2} \equiv 1 \; (\textrm{mod} \; 2^{2k})$ 임을 알 수 있습니다. 한편 소수 $p > 2$에 대해 $(p^{2k}!)_{p} \equiv -1 \; (\textrm{mod} \; p^{2k})$ 인 것은 위키피디아에서 잘 나와있습니다.

 

Hapby9921의 글을 참고하면 $s = (p^{k}!)_{p}$일 때

$$(N!)_{p} \equiv (\pm 1)^{u} \times s^{v} \times \prod_{i=1}^{r} (vp^{k}+i) \;\; \textrm{mod} \; p^{2k}$$

인 것은 큰 무리없이 수용 가능해졌는데, 갑자기 글에서는 위 식을 $N$을 $p$로 나눠가면서 반복하라고 지시합니다. 왜일까요?

 

여기서 분명히 짚고 가야하는 것은, 원래 위 글은 combination을 구하기 위한 알고리즘을 소개하는 것이지, $(X!)_{p}$만을 구하기 위한 알고리즘을 소개한 것이 아닙니다. 즉 $(X!)_{p}$는 위 수식이 이미 구하고 있고, $X$을 $p$로 나눠 반복하는 것이 다른 의미가 있습니다.

 

그리고 그 의미는 Andrew Granville의 글에서 Theorem 1로 등장해버립니다.

$$ \frac{1}{p^{e_{0}}} \binom{n}{m} \equiv (\pm 1)^{e_{q-1}} \prod_{i=0}^{d} \frac{(N_{i}!)_{p}}{(M_{i}!)_{p}(R_{i}!)_{p}} \;\; \textrm{mod} \; p^{q}$$

수열 $\left \{ N_{i} \right \}$의 정의를 읽어보면, $N_{i} = \lfloor \frac{n}{p^{i}} \rfloor$이기 때문에, N을 P로 나눠 반복한다 라는 것이 위 $\prod$ 부분을 계산하기 위함이었다는 걸 짐작할 수 있습니다.

 

이제 Hapby9921의 방법으로 combination을 구하는 것의 시간복잡도를 생각해보겠습니다.

 

우선 $(X!)_{p} \; \% \; p^{2k}$을 구하는 것부터 시작합니다.

 

$s^{v} = (p^{k}!)_{p}^{v}$을 구하는 것은, $(p^{k}!)_{p}$을 별도의 방법 없이 나이브하게 구하기 때문에 $O(p^{k}+\log v)$ 입니다.

 

그리고 $\prod (vp^{k}+i)$ 을 계산하는 것은 자명하게 $O(p^{k})$ 일 것입니다.

 

또한 전체적인 맥락을 봤을 때 $p^{k} \leq D$ 이므로, $(X!)_{p}$을 $O(D)$에 구한다고 볼 수 있겠습니다.

 

그런데 combination을 구하기 위해서는 $(X!)_{p}$의 계산을 $O(\log P)$번 반복해야 합니다. 따라서 대략 $O(D \log D)$ 라고 볼 수 있겠습니다.

 

종만북과 비공식적인 기준(1초당 1억)을 참고했을 때 $D \leq 1.6 \times 10^{7}$이므로 3초 안에 들어오는 지 여부는 제출해봐야 알 수 있습니다. 다만 생략된 복잡도와 상수가 생각보다 영향이 커서 아슬아슬하게 통과됩니다. 저는 register, unsigned type을 포함한 여러 상수 커팅을 시도해봤는데 2300ms 이하로는 줄이기가 매우 힘들었습니다. (C11, 2340ms)

---

이제 rkm0959의 글 로 넘어가겠습니다.

 

이 글에 대해 완벽히 이해한 것도 아니고 정수론을 잘 알지도 못해서, 틀린 점이 있으면 바로 말씀해주시면 감사하겠습니다. 일단 제가 생각하기엔 이 글의 틀린 부분은 $f(m,p,e)$에 관한 점화식입니다.

 

글에서 정수 $m$과 소수 $p$에 대해, $m!=p^{v_{p}(m!)} \times Y$ 일 때 $f(m,p,e) = Y \; \% \; p^{e}$ 라고 정의했는데

$$f(m,p,e) \equiv f(p^{e}-1 , p, e)^{\lfloor \frac{m}{p^{e}} \rfloor} \times f(m \; \% \; p^{e}, p, e) \times f( \lfloor \frac{m}{p^{e}} \rfloor , p, e) \;\; (\textrm{mod} \; p^{e})$$

라는 점화식이 성립한다고 했습니다.

 

예를 들어 $m=13, \; p=e=2$이라면, 언뜻 보기에

$$1,2,3 \;\; / \;\; 4 \;\; / \;\; 5,6,7 \;\; / \;\; 8 \;\; / \;\; 9,10,11 \;\; / \;\; 12 \;\; / \;\; 13$$

와 같이 나누고, 4,8,12와 1,2,3,5,6,7,9,10,11을 잘 처리해주고 마지막에 13을 처리해준다면 위 점화식이 성립하는 것처럼 보입니다.

 

허나, 실제로 계산하게 되면, 먼저 $f(3,2,2)=3$이며, $f(1,2,2)=1$ 이므로, 점화식의 우변은 1입니다.

 

그리고 점화식의 좌변, $f(13,2,2)=3$이므로 글에서 주어진 점화식은 성립하지 않게 됩니다.

 

여기서 Hapby9921과 Granville이 썼던 그대로 $(X!)_{p}$라는 표기를 가져오고, 식을 약간만 바꿔보겠습니다.

$$f(m,p,e) \equiv (p^{e}!)_{p}^{\lfloor \frac{m}{p^{e}} \rfloor} \times ((m \; \% \; p^{e})!)_{p} \times f( \lfloor \frac{m}{p} \rfloor , p, e) \;\; (\textrm{mod} \; p^{e})$$

이 식은 방금 전의 $m=13, \; p=e=2$의 예시에서 잘 verify됩니다. 또한 이 식을 정당화하는 것은 다시 한 번 Granville의 Theorem 1 에서 그 근거를 찾을 수 있습니다. 한편, $e=1$ 일 때 위 두 식은 같습니다.

 

이제 rkm0959의 combination을 구하는 알고리즘의 시간 복잡도를 생각해보겠습니다.

 

첫번째로 $(p^{e}!)_{p}$의 전처리는, 여전히 $O(D)$입니다.

 

하지만 전처리 이후로는 $O(\log N)$입니다. 따라서 전체 시간 복잡도는 roughly $O(D)$ 입니다.

 

그래서 그런지, rkm0959의 알고리즘을 기반으로 한 코드는 수행시간이 9~10배 가량 줄어드는 모습입니다. 다음은 제 코드입니다.

#include <stdio.h>
typedef unsigned long long ull;

int factorize_integer (int X, int primeFactor[10][2]);

ull p_adic_valuation_factorial (ull X, ull P);
void factorize_binomial (ull N, ull M, ull primeFactor[10][2], int numberOfPrime);

ull min (ull A, ull B);

ull power (ull X, ull Y);
ull power_modular (ull X, ull Y, ull M);

ull factorial_CoprimeToP_moduloPrimePower_initial[16000100];
ull factorial_CoprimeToP_moduloPrimePower_general (ull N, ull M, ull P);
ull binomial_residue_moduloPrimePower (ull N, ull M, ull primeFactor[10][2], int numberOfPrime, ull P, ull Q);

ull Chinese_Remainder_Theorem (ull result[10], ull modulus, int primeFactor[10][2], int numberOfPrime);

int main(void) {
	
	ull N, M, D;
	scanf("%llu %llu %llu", &N, &M, &D);
	
	int primeFactorize_D[10][2] = {{}};
	int numberOfPrimeFactor_D = factorize_integer(D, primeFactorize_D);
	
	ull primeFactorize_NchooseM[10][2] = {{}};
	for (int i = 0; i < numberOfPrimeFactor_D; i++) {
		primeFactorize_NchooseM[i][0] = primeFactorize_D[i][0];
	}
	factorize_binomial(N, M, primeFactorize_NchooseM, numberOfPrimeFactor_D);
	
    ull powerOfD_in_NchooseM = primeFactorize_NchooseM[0][1] / primeFactorize_D[0][1];
	for (int i = 0; i < numberOfPrimeFactor_D; i++) {
		powerOfD_in_NchooseM = min(powerOfD_in_NchooseM,
		                           primeFactorize_NchooseM[i][1] / primeFactorize_D[i][1]);
	}
	
	ull result[10] = {};
	for (int i = 0; i < numberOfPrimeFactor_D; i++) {
		result[i] = binomial_residue_moduloPrimePower(N, M,
													  primeFactorize_NchooseM, numberOfPrimeFactor_D,
													  primeFactorize_D[i][0], primeFactorize_D[i][1]);
	}
	
	ull answer = Chinese_Remainder_Theorem(result, D, primeFactorize_D, numberOfPrimeFactor_D);

	for (int i = 0; i < numberOfPrimeFactor_D; i++) {
		answer = answer * power_modular(primeFactorize_NchooseM[i][0],
		                                primeFactorize_NchooseM[i][1] - primeFactorize_D[i][1] * powerOfD_in_NchooseM,
		                                D)
		                % D;
	}

	printf("%llu", answer);
	
	return 0;
}

int factorize_integer (int X, int primeFactor[10][2]) {
	// get [(p_i, a_i)] for an integer X <= 1.6 * 10^7
	
	int numberOfPrimeFactor = 0;
	if (X % 2 == 0) {
		primeFactor[numberOfPrimeFactor][0] = 2;
		while(X % 2 == 0) {
			X /= 2;
			primeFactor[numberOfPrimeFactor][1]++;
		}
		numberOfPrimeFactor++;
	}
	for (int i = 3; i * i <= X; i += 2) {
		if (X % i == 0) {
			primeFactor[numberOfPrimeFactor][0] = i;
			while(X % i == 0) {
				X /= i;
				primeFactor[numberOfPrimeFactor][1]++;
			}
			numberOfPrimeFactor++;
		}
	}
	if (X > 1) {
		primeFactor[numberOfPrimeFactor][0] = X;
		primeFactor[numberOfPrimeFactor][1] = 1;
		numberOfPrimeFactor++;
	}
	return numberOfPrimeFactor;
}

ull p_adic_valuation_factorial (ull X, ull P) {
	// get the exponent of the highest power of prime P that divides X!
	// Legendre, 1830. v_p(X!) = sum floor(X / p^i) (1 <= i)
	
	ull result = 0;
	for (X /= P; X; X /= P) {
		result += X;
	}
	return result;
}

void factorize_binomial (ull N, ull M, ull primeFactor[10][2], int numberOfPrime) {
	// get [(p_i,a_i)] where (N choose M) = p1^a1 * p2^a2 * ... pn^an * Residue
	// p_i are given in [i][0]; put a_i in [i][1]
	
	ull R = N - M;
	for (int i = 0; i < numberOfPrime; i++) {
		primeFactor[i][1] = p_adic_valuation_factorial(N, primeFactor[i][0]) -
		                    p_adic_valuation_factorial(M, primeFactor[i][0]) -
		                    p_adic_valuation_factorial(R, primeFactor[i][0]);
	}
}

ull min (ull A, ull B) {
	return A > B ? B : A;
}

ull power (ull X, ull Y) {
	ull result = 1;
	while (Y) {
		if (Y % 2) {
			result = result * X;
		}
		X = X * X;
		Y /= 2;
	}
	return result;
}

ull power_modular (ull X, ull Y, ull M) {
    ull result = 1;
	while (Y) {
		if (Y % 2) {
			result = result * X % M;
		}
		X = X * X % M;
		Y /= 2;
	}
	return result;
}

ull factorial_CoprimeToP_moduloPrimePower_general (ull N, ull M, ull P) {
	ull result = 1;
	while(N) {
		result = result * (factorial_CoprimeToP_moduloPrimePower_initial[N % M] * power_modular(factorial_CoprimeToP_moduloPrimePower_initial[M-1], N/M, M) % M) % M;
		N /= P;
	}
	return result;
}

ull binomial_residue_moduloPrimePower (ull N, ull M, ull primeFactor[10][2], int numberOfPrime, ull P, ull Q) {
	// get Residue % P^Q where N choose M = p1^a1 * p2^a2 * ... * pn^an * Residue
	
	ull modulus = power(P, Q);
	
	factorial_CoprimeToP_moduloPrimePower_initial[0] = 1;
	for (ull temp = 1, i = 1, j = 1; i < modulus; i++, j++) {
		if (j == P) j = 0;
        else temp = temp * i % modulus;
		factorial_CoprimeToP_moduloPrimePower_initial[i] = temp;
	}
	
	ull R = N - M;
	ull phiModulus = modulus - modulus / P;
	
	ull answer = factorial_CoprimeToP_moduloPrimePower_general(N, modulus, P);
	answer = answer * power_modular(factorial_CoprimeToP_moduloPrimePower_general(M, modulus, P) *
                                    factorial_CoprimeToP_moduloPrimePower_general(R, modulus, P) % modulus, phiModulus - 1, modulus) % modulus;
	for (int i = 0; i < numberOfPrime; i++) {
		if (primeFactor[i][0] == P) continue;
		answer = answer * power_modular(primeFactor[i][0], primeFactor[i][1] * (phiModulus - 1), modulus) % modulus;
	}
	
	return answer;
}

ull Chinese_Remainder_Theorem (ull result[10], ull modulus, int primeFactor[10][2], int numberOfPrime) {
	
	ull temp[10][4] = {{}};
	/*
	 * temp[*][0] = m_i = (p_i ^ a_i) in D
	 * temp[*][1] = phi(temp[*][0])
	 * temp[*][2] = M_i = modulus / m_i
	 * temp[*][3] = inverse(temp[*][2]) (mod temp[*][0])
	 * where ans = sum [*][2] * [*][3] * result[*] (Chinese Remainder Theorem)
	*/
	for (int i = 0; i < numberOfPrime; i++) {
		temp[i][0] = power(primeFactor[i][0], primeFactor[i][1]);
		temp[i][1] = temp[i][0] - temp[i][0] / primeFactor[i][0];
		temp[i][2] = modulus / temp[i][0];
		temp[i][3] = power_modular(temp[i][2], temp[i][1] - 1, temp[i][0]);
	}
	
	ull answer = 0;
	for (int i = 0; i < numberOfPrime; i++) {
		answer = (answer + result[i] * (temp[i][2] * temp[i][3] % modulus)) % modulus;
	}
	return answer;
}

(C11, 126116KB, 244ms, 제출번호 56432971)

---

Granville의 글을 적용한 코드는 구현해보지 않았지만, 아마 성능으로는 Hapby9921보다 안 좋게 나오지 않을까 예상됩니다. 왜냐하면, $O(\log ^{2} N + Q^{4}\log N \log P + Q^{4}P\log ^{3}P)$ 인데, $P \approx 1.6 \times 10^{7}$이라면 $O(D\log ^{3}D)$에 가까워질 거라고 예상하기 때문입니다. 물론 큰 $P$에 대해 Lucas' Theorem을 쓰면 다시 $O(D)$에 가까워질 수도 있겠지만, 성능이 얼마나 나올 지는 구현을 해봐야 알 것 같습니다.

---

이 문제를 통해 알게 된 유용한 최적화는, $P$의 배수인지 아닌지, 혹은 $P$번째 반복 중인지를 더 빠르게 돌리는 것입니다. 즉

for (int i = 1; i < N; i++) {
    if (i % P == 0) {
    
    }
}

보다

for (int i = 1, j = 1; i < N; i++, j++) {
    if (j == P) {
    	j = 0;
    }
    else {
    
    }
}

로 처리하는 게 더 빨랐습니다.