[BOJ/백준 11401번][PYTHON] 이항계수 3

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

 

이런 문제가 젤 무섭다. 문제가 한 줄로 끝나는 데 골드 1에 있고, 시간 제한 1초, 메모리 제한 256MB로 딱봐도 타이트하게 걸려있다.

또한 1,000,000,007로 나눈 나머지라는 말이 있는 거로 봐선 입력에도 큰 값이 들어올 것이라는 것을 당당히 밝히고 있다.

# 1,000,000,007인 이유 :  (int의 범위 : -2147483648 ~ 2147483647) int의 max에 가까운 수 중에서 절반에 해당하는 소수 값

1. 모듈러 연산을 하기 위해서

2. stackoverflow 방지 (최대값으로 나눌 경우, overflow일어날 수 있음)

 

이제 해당 문제에 대한 접근을 하게 되면, 보통의 이항 정리 공식이라고 하면

nCr => n! / ((n-r)! (r)!) 이렇게 된다.

이렇게 해결될 문제면 브론즈에 있지 골드에 없을 것이다.

 

해당 방식을 통해서 단순 구현한 후 돌리게 되면,

ㅋㅋㅋㅋㅋㅋ

이 중 하나가 되게 된다. 

 

그리하여 어떤 방법이 있나 하고 생각을 해보았다. 분명 수학적 지식이 들어간다 생각하고 배운 부분 중에 쓸만한 것이 있나 생각해봤다.

근데 아니다 다를까 아무 생각이 안났다. (아는 후임한테 물어봤더니 수학 지식 필요한 문제는 어쩔 수 없다는 대답이 왔다.)

 

알고리즘 분류에 이런 것이 보였다.

페르마 소정리?

페르마의 소정리란, a^p ≡ a (mod p) 이렇게 된다. 그렇게 되면 a^(p-1) ≡ 1 (mod p),

a^(p-2) ≡ 1/a 이렇게 된다. -> 이 부분이 중요

 

아니 이거로 어떻게 하지? (블로그를 참조함)

 

다시 식을 가져와보자 n! / (r!(n-r)!)

nCr % p ≡ n!/((n-r)!(r)!) % p ≡ n!((n-k)!(k)!)^-1 % p 

여기서 (k! * (n-k)!) -> (r! * (n-r)!)^(p-2)로 바꿨다.  (이 부분이 페르마 소정리를 이용한 부분)

 

하지만 단순히 이 부분만 바꾸고 구현을 하게 되면, 재귀 오류나 시간 오류를 피하기 힘들다. 따라서 이 부분에 대해서도

생각을 하고 코드 작성이 필요하다.

해당 문제는 분할 정복의 문제 중 하나였다. 따라서 분할 정복의 아이디어를 활용할 필요가 있다.

def divide_fac(n,m):
    if n == m:
        return n
    if m < n:
        return 1
    else:
        return divide_fac(n , (n+m)//2) * divide_fac((n+m)//2+1,m) % prime_number

# 큰 수의 곱셈을 작은 단위로 나눠서 계산하였다.

tmp = divide_fac(1,n)    # n!
tmp1 = divide_fac(1,n-k) # (n-k)!
tmp2 = divide_fac(1,k)   # k!

# 이렇게 3가지 수를 각각 구하였다.

def divide_pow(a, x):
    if x == 0:
        return 1
    if x == 1:
        return a % prime_number
    
    half = divide_pow(a, x//2) % prime_number
    
    if x % 2 == 0:
        return (half * half) % prime_number
    else:
        return (half * half * a) % prime_number

# 이 부분의 거듭제곱 또한 분할 정복을 통하여 구하였다. 이 부분이 갑자기 떠오르지 않아서 오래 걸렸다.

# x가 짝수인 경우와 훌수인 경우를 나눠서 연산을 진행하였다. a가 홀수라면 정확히 절반을 나누기 힘들기 때문에

# 차이를 두어 계산하였다.

 

전체 코드이다.

import sys
n, k = map(int, sys.stdin.readline().split())

prime_number = 1000000007
dp = [0] * (n+1)
dp[1] = 1

def divide_fac(n,m):
    if n == m:
        return n
    if m < n:
        return 1
    else:
        return divide_fac(n , (n+m)//2) * divide_fac((n+m)//2+1,m) % prime_number

tmp = divide_fac(1,n) # 5, n은 5이고 k는 2
tmp1 = divide_fac(1,n-k) # 3
tmp2 = divide_fac(1,k) # 2

def divide_pow(a, x):
    if x == 0:
        return 1
    if x == 1:
        return a % prime_number
    
    half = divide_pow(a, x//2) % prime_number
    
    if x % 2 == 0:
        return (half * half) % prime_number
    else:
        return (half * half * a) % prime_number

print((tmp*divide_pow(tmp1 * tmp2, prime_number-2))%prime_number)

 

얻어 간 점은 시간을 줄이기 위해서 정수론의 개념이 사용가능하며, 그 아이디어가 페르마 소정리에서 기인한다는 것이다.