[파이썬] 백준 11401번: 이항 계수 3 (분할 정복, 페르마의 소정리)

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

11401번: 이항 계수 3

 

문제

자연수 N\(N\)과 정수 K\(K\)가 주어졌을 때 이항 계수 (NK)\(\binom{N}{K}\)를 1,000,000,007로 나눈 나머지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N\(N\)과 K\(K\)가 주어진다. (1 ≤ N\(N\) ≤ 4,000,000, 0 ≤ K\(K\) ≤ N\(N\))

출력

 (NK)\(\binom{N}{K}\)를 1,000,000,007로 나눈 나머지를 출력한다.

예제 입력 1 복사

5 2

예제 출력 1 복사

10

 

N의 범위가 1000000이상으로 큰 경우 그냥 계산하면 시간초과가 남

모듈러 연산을 통한 분할 정복이 필요 + 페르마의 소정리

# 조합 + 모듈러 연산 + 페르마 소정리 + 분할 정복
import sys
input=sys.stdin.readline

def power(a, b):
    if b == 0:
        return 1
    if b % 2:   #홀수이면
        return (power(a, b//2) ** 2 * a) % p
    else:
        return (power(a, b//2) ** 2) % p

p = 1000000007
N, K = map(int, input().split())

fact = [1 for _ in range(N+K+1)]

for i in range(2, N+1):
    fact[i] = fact[i-1] * i % p # 팩토리얼 미리 계산해둠

# 조합 공식
A = fact[N] # 분자
B = (fact[N-K] * fact[K]) % p #분모

# 페르마의 소정리에 의해 a^(p−2)=(1/a)mod p 임
# 즉 분모의 경우에는 a^(p−2)로 변환이 필요 
print((A % p) * (power(B, p-2) % p) % p )