BJ11439 이항 계수 5

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

정공법이 잘 먹히는 문제입니다.

소수가 아닌 M

소수가 아닌 $m$에 대하여

을 구하는 가장 좋은 방법은 직접 해보는 겁니다.

이상한 답변이긴 하지만, $m$을 소인수분해하여 $p_1^{e_1}p_2^{e_2}\ldots p_k^{e_k}$로 만들어 각각을 계산한다고 하더라도, 지수의 $e_i$가 문제가 됩니다(이는 이후 이항 계수 6을 어렵게 만드는 원인입니다).

그 대신, $n!$, $(n-k)!$, $k!$을 빠른 시간에 소인수분해 할 수 있으면 가능하지 않을까요?

Legendre 공식

르장드르 공식(Legendre's formula)에 따르면, $n!$의 소인수 분해 $p$의 지수 $v_p(n!)$는 다음과 같습니다.

그렇다면 $n$ 이하의 소수 $p$에 대해 각각의 $v_p(n!)$을 구하는 데 필요한 연산의 수는 $\log_p n$이고, $n$ 이하의 소수의 개수는 대략 $n/\log n$이므로 소수들의 평균을 $n/2$로 가정하면 총 연산수는

라고 생각하면 되겠네요. $n=4000000$에 약 28만 정도이니, 크게 문제없을 것 같네요.

에라토스테네스의 체

이제 $n$ 이하의 소수를 구해야 합니다.

2와 3을 먼저 소수 리스트에 추가하고, 5부터 시작해 봅시다. 2보다 큰 소수는 항상 홀수이니 2씩 건너 뛰어가며 $n$ 이하의 홀수를 소수 후보로 둡시다. 만일 소수 리스트의 모든 수에 대해 소수 후보가 나누어 떨어지지 않는다면 정의에 따라 소수 후보는 소수입니다.

또한 소수 리스트의 모든 수를 체크하는 것이 아닌, $\sqrt n$ 이하의 수만 체크해도 됩니다. 다만 이를 실제로 구현하면 소수가 $\sqrt n$ 이하인지 체크하는 부하가 생길 것입니다. 환경에 따라 어느 것의 성능이 우위라고 할지 결정할 수는 없습니다.

코드 - 첫 번째 시도

n, k, m = map(int, input().split())
nk = n-k

primes = [2, 3]
for prime_candidate in range(5, n+1, 2):
    flag = False
    for prime in primes:
        if prime**2 > prime_candidate:
            break
        if prime_candidate % prime == 0:
            flag = True
            break
    
    if not flag:
        primes.append(prime_candidate)

n_dec, nk_dec, k_dec = dict(), dict(), dict()
def calc_dec(num: int, dec: dict):
    for prime in primes:
        prime_pow = prime
        while prime_pow <= num:
            dec[prime] = dec.get(prime, 0) + num // prime_pow
            prime_pow *= prime
calc_dec(n, n_dec)
calc_dec(nk, nk_dec)
calc_dec(k, k_dec)
        
binom_dec = dict()
for key, value in n_dec.items():
    exp = n_dec[key] - nk_dec.get(key, 0) - k_dec.get(key, 0)
    if exp == 0:
        continue
    else:
        binom_dec[key] = exp

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

result = 1
for key, value in binom_dec.items():
    result *= pow(key, value, m)
    result %= m

print(result)

실행 결과 Python3을 기준으로는 시간 초과, PyPy로만 통과하나 그 속도가 느립니다.

개선할 부분이 있을까요?

primes 리스트의 append()?

코드 수행 시간을 측정해 보면 병목 현상이 나는 코드는 에라토스테네스의 체 부분입니다.

에라토스테네스 구현 자체를 바꾸는 것은 어려워 보이고, 시간복잡도 자체가 $O(n\log(\log n))$이므로 수행 시간 내에는 해결할 수 있습니다. 어쩌면 append가 성능 저하를 발생시키는 요인일 수도 있습니다.

소수 계량 함수에 따르면, $\pi(x)\approx x/\log x$입니다. 큰 $x$에 대해 $\pi(x)>x/\log x$이고, $x=4000000$일 때 각각이 약 30,000 가량 차이가 납니다. 따라서 primes 배열을 다음과 같이 정의할 수 있습니다.

primes = [0] * (int(n / log(n)) + 30000)
primes[0] = 2
primes[1] = 3

또한 이 때 primes 배열을 만드는 방법은

counter = 2
for prime_candidate in range(5, n+1, 2):
    flag = False
    for prime in primes:
        if prime**2 > prime_candidate:
            break
        if prime_candidate % prime == 0:
            flag = True
            break
    
    if not flag:
        primes[counter] = prime_candidate
        counter += 1

이 될 겁니다. 그리고 calc_dec에서도 primes 배열은 primes[:counter]로 접근해야 할 것입니다.

결과
이전 것은 1992 ms였는데, 2044 ms로 오히려 늘어났네요.
이 방법은 아닌 것 같습니다.

built-in 클래스, 함수 이용

Pythonic한 에라토스테네스의 방법이 있어 소개합니다.

https://velog.io/@joygoround/test-unsolved-%EC%86%8C%EC%88%98-%EC%B0%BE%EA%B8%B0-%ED%8C%8C%EC%9D%B4%EC%8D%AC

어떤 길이 n의 generator로부터 set을 생성하는 시간복잡도는 O(len(...))이고, 원소가 각각 s개, t개인 두 set의 차집합을 구하는 시간복잡도는 O(len(s)+len(t)) 입니다. Built-in 메서드이니 Python 내에서의 최적화도 기대해 볼 수 있겠습니다. 다만, 어차피 3 이상의 소수는 항상 홀수이니 다음과 같이 최적화할 수 있겠습니다.

primes = set(range(3, n+1, 2))
primes.add(2)
for num in range(3, n+1, 2):
    if num in primes:
        primes -= set(range(2*num, n+1, num))

다만 메모리 사용량이 늘어나는 것은 감안해야합니다. 시간도 1860 ms로 크게 줄지 않은 것을 생각하면, 첫번째 방법이 더 나은 것으로 사료됩니다.