[Python] 백준 15718 - 돌아온 떡파이어

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

문제

떡파이어의 불로장생의 비밀은 바로 떡국이다.

떡파이어는 떡국을 먹은 그릇의 개수만큼 나이를 먹는다. 그들은 매일 떡국을 먹는데, 떡국을 먹는대로 바로 소화가 가능하기 때문에 하루에 얼마든지 원하는 만큼 떡국을 먹을 수 있다. 그러나 전에 떡국을 얼마나 먹었든지, 그들은 기구하게도 떡국을 하루라도 먹지 않으면 생을 마감하게 된다.

어느 날, 디디는 어떤 떡파이어가 M째날에 N세로 생을 마감하기까지 어떤 생을 살아왔는지 알고 싶어서, 그의 나이를 먹는 과정의 경우의 수를 세려고 한다. 그렇지만, 떡파이어의 나이가 많을 수록 그 경우의 수는 무수히 많아지기 때문에 디디는 곤란해하고 있다.

그런 디디를 위해 M째날에 N세로 생을 마감한 떡파이어가 나이를 먹는 과정의 경우의 수를 세는 프로그램을 작성해야 한다.

떡파이어의 나이는 0세부터 시작된다.

N = 3, M = 3,일때를 예로 들면,

• 첫째 날 1개 둘째 날 2개, 셋째 날 0개
• 첫째 날 2개 둘째 날 1개, 셋째 날 0개

총 경우의 수는 2이다.

입력

첫째 줄에 테스트 케이스의 수 T(1 ≤ T ≤ 1000)가 주어진다.

각 테스트 케이스는 한 줄로 이루어져 있으며, 정수 N(0 ≤ N ≤ 10^9)과 M(1 ≤ M ≤ 10^9)이 공백으로 구분되어 주어진다.

출력

각 테스트 케이스마다 한 줄에 하나씩 나이를 먹는 방법의 가짓 수를 100007로 나눈 나머지를 출력하시오. 100007은 일반적이지 않은 나눗수임에 유의하라.

 

---

풀이

이 문제를 풀려면 두 가지 수학적인 개념을 알고 있어야 한다:

 

1. 뤼카의 정리

2. 중국인의 나머지 정리

 

이 두 가지 개념에 대해서는 추후에 블로그에 정리해서 올려볼 예정이다.

 

문제를 살펴보면 문제의 내용은 아주 전형적인 확률과 통계 조합론 문제임을 알 수 있다.

M일동안 N개의 떡국을 먹어야하는데, 일단 마지막날은 떡파이어가 죽어야하므로 마지막날 먹은 떡국의 개수는 0개로 고정이다.

그리고 살아있을 때에는 최소 떡국을 1그릇은 먹어야하므로, 문제를 이렇게 재정의할 수 있겠다:

 

자연수 M-1개를 더해서 N이 나오게 하는 경우의 수는 몇 가지인가?

 

결론부터 말하자면 n-1 C m-2 를 구하면 된다. 그 이유는, 

1을 n개 늘어놓았다고 생각해보자. 

1, 1, 1, 1, 1, 1, ........ ,1

이제 m-2개의 칸막이를 이용하여 이를 m-1개의 칸으로 나누면 된다. 각 칸에 있는 1의 개수가 그 날 먹은 떡국의 개수일것이다.

칸막이를 넣을 수 있는 공간이 n-1개 있으므로 n-1Cm-2를 구하면 된다.

 

단, n과 m의 범위가 10^9까지나 되므로, 단순히 math.comb를 쓸 것이 아니라 뤼카의 정리를 이용해주어야 할 것이다.

 

아래의 코드는 내가 예전에 작성한 뤼카의 정리 코드이다.

 

def fac(n, P):
    ans = 1
    for i in range(1, n+1):
        ans = (ans*i)%P
    return ans

def com(n, k, P):
    if n < k:
        return 0
    num = (fac(n-k, P) * fac(k, P))%P
    large = pow(num, P - 2, P)
    return (fac(n, P) * large) % P

def C(n, k, P):
    i = 0
    while P ** i <= n:
        i += 1
    i -= 1
    arr1 = []
    arr2 = []

    while True:
        if i == -1:
            break
        arr1.append(n // (P ** i))
        n %= (P ** i)
        arr2.append(k // (P ** i))
        k %= (P ** i)
        i -= 1

    final = 1
    for i in range(len(arr1)):
        final *= com(arr1[i], arr2[i], P)
        final %= P
    return final

(사실 예전에 작성한거라 좀 더럽고 길다)

(어쨌든 작동은 잘 한다)

 

그런데, 뤼카의 정리는 nCr의 값을 소수 p로 나눈 나머지를 구하는 방법이다.

문제에서는 100007로 나눈 나머지를 구하라고 했다. 

 

함정카드

 

그래서 혹시 몰라서 100007을 소인수분해 해보았다.

 

 

100007 = 97*1031

 

 

100007은 소수가 아니기때문에 뤼카의 정리 코드에 p 대신 100007를 넣게되면 백퍼 WA를 받는다.

여기서 필요한 것이 중국인의 나머지 정리이다.

 

법 97에 대한 해와 법 1031에 대한 해가 있을 때, 법 100007에 대한 해는 반드시 딱 하나만 존재한다는 것이 중국인의 나머지 정리이다.

 

중국인의 나머지 정리의 해를 구하려면

 

a1 = n-1Cm-2 % 1031

a2 = n-1Cm-2 % 97

 

을 뤼카의 정리를 이용하여 얻는다.

 

그 후 p1 = 1031, p2 = 97을 구한 다음

각각을 p2, p1의 모듈로 역원을 구해야 한다

 

이를 m1, m2라고 하면 페르마의 소정리를 이용하여(97, 1031은 소수이므로) 역원을 쉽게 구할 수 있다.

럭키비키하게도 파이썬에서는 a^b (mod p)를 굉장히 빠르고 쉽게 pow()를 이용하여 구할 수 있다

 

a1*p2*m1 + a2*p1*m2 가 100007에 대한 최종적인 해가 될 것이다.

따라서 중국인의 나머지 정리 코드는 다음과 같다:

 

a = C(N-1, M-2, 1031)
b = C(N-1, M-2, 97)
m1 = pow(97, 1029, 1031)
m2 = pow(1031%97, 95, 97)
ans = a*97*m1 + b*1031*m2
print(ans%100007)

 

 

하지만 이 문제에는 함정이 3개나 더 있다. 엣지 케이스가 살짝 까다롭기 때문이다.

 

1. N == 0일때

 

N이 0이면 하루만에 태어나서 하루만에 죽어야한다. 즉, M이 1인 경우를 제외하고는 전부 답이 0이다.

 

2. M == 1일때

 

1번과 비슷하게, 살 날이 하루 밖에 없다는 뜻이므로 N이 0인 경우를 제외하고는 전부 답이 0이다.

(N, M) = (2, 1)과 같은 경우때문에 이번 단계에서 한 번 더 걸러줄 필요가 있다.

 

3. M-1 > N일때(이걸 못찾아서 하루종일 고생했다)

하루 안에 최소 1개의 떡국은 먹어야하므로 죽기 전까지 최소한 M-1개의 떡국을 먹어야한다. 그런데 먹을 날은 M-1개 주어지는데 이게 N보다 크면 비둘기집의 원리에 의해 무조건 1개의 날은 떡국 0개를 먹어야하므로 말이 안된다.

 

그래서 이 3개의 경우를 모두 찾아서 특수처리 해주고, 그 이외의 경우에 대해서는 뤼카의 정리와 중국인의 나머지 정리를 이용해서 답을 내면 AC를 받을 수 있다.

 

# 100007 = 97*1031
def fac(n, P):
    ans = 1
    for i in range(1, n+1):
        ans = (ans*i)%P
    return ans

def com(n, k, P):
    if n < k:
        return 0
    num = (fac(n-k, P) * fac(k, P))%P
    large = pow(num, P - 2, P)
    return (fac(n, P) * large) % P

def C(n, k, P):
    i = 0
    while P ** i <= n:
        i += 1
    i -= 1
    arr1 = []
    arr2 = []

    while True:
        if i == -1:
            break
        arr1.append(n // (P ** i))
        n %= (P ** i)
        arr2.append(k // (P ** i))
        k %= (P ** i)
        i -= 1

    final = 1
    for i in range(len(arr1)):
        final *= com(arr1[i], arr2[i], P)
        final %= P
    return final


T = int(input())
for _ in range(T):
    N, M = map(int, input().split())
    if N == 0:
        if M == 1:
            print(1)
        else:
            print(0)

    elif M == 1:
        if N != 0:
            print(0)
        else:
            print(1)

    elif  M-1>N:
        print(0)

    else:
        a = C(N-1, M-2, 1031)
        b = C(N-1, M-2, 97)
        m1 = pow(97, 1029, 1031)
        m2 = pow(1031%97, 95, 97)
        ans = a*97*m1 + b*1031*m2
        print(ans%100007)

최종적인 코드이다.

나중에 뤼카의 정리 코드를 살짝 더 줄이는 작업을 거쳐봐야겠다.