[BOJ 11401] 이항 계수 3

💯 코딩테스트/[2021~] 💯 코딩테스트

[BOJ 11401] 이항 계수 3 - Python 풀이

락꿈사 2022. 2. 7. 09:10

문제

입력

출력

코드

import sys
P = 1000000007

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

def solve(N, K):
    fact = [1 for _ in range(N+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

    print((A % P) * (power(B, P-2) % P) % P )

if __name__ == "__main__":
    N, K = map(int, sys.stdin.readline().split())
    solve(N,K)

결과

간단한 문제이나 N의 크기가 40,000,000이나 되는게 함정인 문제였다.

n!/(n-k)k! 을 이용해서 풀려고 했으나, 1,000,000,007로 나눠야 하므로 분자와 분모에 나머지를 구해서 계산하게 되면 0으로 나누는 문제가 발생할 수 있기 때문에 분수꼴을 곱셈꼴로 바꿔야 한다.

(https://rhdtka21.tistory.com/123 <- 링크 참조)

따라서 페르마의 소정리를 이용해야 한다

페르마의 소정리는 p가 소수이고 a가 p의 배수가 아니면