[BOJ, Python] 11401 - 이항 계수 3 ( with Fermat's little theorem )

11401번: 이항 계수 3

PS의 난이도가 올라가다 보면 같은 문제를 좀 더 효율적인 시간, 공간 복잡도로 구현해야 하는 순간이 온다. 이항 계수를 구할 때 n의 값이 너무 크면 메모이제이션(dp)으로도 해결하지 못한다. 문제가 요구하는 정답이 이항 계수의 모듈러 연산 값이라면, 페르마의 소정리와 곱셈의 역원을 활용하여 답을 찾을 수 있겠다.

 

 

페르마의 소정리

p가 소수라면,

$$ a^{p - 1} \equiv 1\pmod{p} $$

 

해당 식의 좌변과 우변이 동치이므로,

 

 

$$\therefore a \times a^{p - 2} \equiv 1 \pmod {p} $$

 

이항 계수 식에 적용

 

 

모듈러 연산을 활용하면 두 수의 곱의 나머지 연산을 간단히 할 수 있다는 장점이 있다.

그런데 이를 이항 계수의 곱에 적용하기엔 무리가 있어보인다.

 

우리는 이항 계수를 의미하는

$$ _{n}\mathrm {C}_{k} = \frac {n!}{k!(n-k)!} $$

이 식에 모듈러 연산을 적용하려면, 분모에 있는

$$ \frac {1}{k!(n-k)!} $$

를 정수의 모듈러 연산으로 환원시켜야 하고, 이때 곱셈의 역원을 사용한다.

 

$$ \frac {1}{a}  \equiv a^{p - 2}\pmod {p} $$

임을 활용하면

 

 

$$ _{n}\mathrm {C}_{k}\bmod p = n!\bmod p \times n!(n-k)!^{p - 2}\bmod p $$

로 표현할 수 있겠다.

 

소스 코드

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

# 거듭 제곱을 할 때, 모듈러 연산을 해준다.
def power(x, y):
    ans = 1
    while y > 0:
        if y % 2 == 1:
            ans *= x
            ans %= m
        x *= x
        x %= m
        y = y // 2
    return ans

# 팩토리얼과 역원을 담는 리스트
fact = [0 for _ in range(n + 1)]
inv = [0 for _ in range(n + 1)]

# 팩토리얼 0, 1 초기화
fact[0] = 1
fact[1] = 1

# bottom up 방식으로 팩토리얼 전처리
for i in range(1, n + 1):
    fact[i] = (fact[i - 1] * i) % m

# top down 방식으로 역원 전처리
inv[n] = power(fact[n], m - 2)

for i in range(n - 1, -1, -1):
    inv[i] = (inv[i + 1] * (i + 1)) % m

# 최종 계산
if n == k or k == 0:
    print(1)
else:
	# (n - k)!의 역원 계산(+ 모듈러 연산)
    r = (fact[n] * inv[n - k]) % m
    # k!의 역원 계산(+ 모듈러 연산)
    r = (r * inv[k]) % m
    print(r)