[440] 조합(combination) 큰 수 계산하기(모듈러) - 파이썬

1. 조합 계산 코드(파이썬)

n = 1000000
r = 500000
MOD = 1000000007

facto = [1]*(n+1)
inv_facto = [1]*(n+1)

for i in range(2, n+1):
    facto[i] = facto[i-1] * i % MOD

inv_facto[n] = pow(facto[n], MOD-2, MOD) # 페르마의 소정리
for i in range(n-1, -1, -1):
    inv_facto[i] = inv_facto[i+1] * (i+1) % MOD

def nCr(n, r):
    return facto[n] * inv_facto[r] % MOD * inv_facto[n-r] % MOD

print(nCr(5, 3))
# 10

print(nCr(10, 4))
# 210

print(nCr(100, 4))
# 3921225

print(nCr(1000000, 500000))
# 996692777

2. math 모듈의 comb 함수와 비교

from time import time
from math import comb
s = time() # 시작 시간

n = 1000000
r = 500000
MOD = 1000000007
facto = [1]*(n+1)
inv_facto = [1]*(n+1)
for i in range(2, n+1):
    facto[i] = facto[i-1] * i % MOD
inv_facto[n] = pow(facto[n], MOD-2, MOD) # 페르마의 소정리
for i in range(n-1, -1, -1):
    inv_facto[i] = inv_facto[i+1] * (i+1) % MOD
def nCr(n, r):
    return facto[n] * inv_facto[r] % MOD * inv_facto[n-r] % MOD
print(nCr(1000000, 500000))
# 996692777

e = time() # 종료 시간
print(f"{e-s:.5f}")
# 0.18896

s = time() # 시작 시간

print(comb(1000000, 500000)%MOD)
# 996692777

e = time() # 종료 시간
print(f"{e-s:.5f}")
# 8.57130

같은 값을 출력하지만 시간 차이가 많이 나는 것을 확인할 수 있습니다.

아마 comb 함수는 내부적으로 n!을 (n-r)!로 나누는 것으로 보입니다.

연산 과정 사이사이에 모듈러 연산을 하지 않아 큰 수를 연산하는 과정에서 시간이 오래 소요되는 것으로 보입니다.

3. 원리

팩토리얼과 팩토리얼의 곱셈에 대한 역원을 미리 계산해두고(동적 계획법), 필요한 조합(combination) 계산에서 값만 가져다 연산함으로써 O(1)에 연산이 가능합니다.

이것은 조합 계산을 많이 해야하는 경우에 더더욱 효율적입니다.

 

팩토리얼을 계산하는 것은 쉬우나, 계속해서 모듈러 연산을 통해 값의 크기가 일정 값 이상으로 커지지 않습니다.

 

다음으로 곱셈에 대한 역원을 계산합니다.

방법은 n!을 계산한 후, 페르마의 소정리로 n!의 곱셈에 대한 역원을 계산하는 것부터 시작됩니다.

이어서 (n-1)!의 곱셈에 대한 역원은 n!의 곱셈에 대한 역원에 n을 곱해서 바로 구할 수 있고,

이를 통해 1!의 곱셈에 대한 역원까지 O(n)으로 구할 수 있습니다.

 

\(_n C _r = \dfrac{_n P _r}{r!} = \dfrac{n!}{r!(n-r)!}\)

조합 계산 방법으로 저장된 값을 활용하여 구하면 연산이 끝납니다.