BOJ 11402 - 이항 계수 4 (Python3)

사실 이건 제법 쟁여두려고 했으나...

작성자 본인이 처음 푼 플래티넘 V 이상 문제이기도 하고, 며칠 간 고민하다 겨우 푼 문제라 그런가 감회가 새롭기도 하다

게다가 이미 방법론은 다른 경로로 알고 있어서 더욱. 그래서 오늘의 문제로 11402번 문제를 택하기도 했다.

열화판으로 11401번 문제가 있는데, 접근 방식은 아예 다르니 조심하길.

 

잡설이 길었다! 그럼 본격적으로 풀이에 돌입해 보자.

---

문제

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

 

입력

첫째 줄에 N, K와 M이 주어진다. (1 ≤ N ≤ \(10^{18}\), 0 ≤ K ≤ N, 2 ≤ M ≤ 2000, M은 소수)

 

예제 입력)

# 예제 입력 1
5 2 3

# 예제 입력 2
30 10 3

# 예제 입력 3
30 3 3

# 예제 입력 4
100 45 7

# 예제 입력 5
100 45 13

 

출력

\(\binom{N}{K}\)를 M으로 나눈 나머지를 출력한다.

 

예제 출력)

# 예제 출력 1
1

# 예제 출력 2
0

# 예제 출력 3
1

# 예제 출력 4
0

# 예제 출력 5
2

 

---

내 코드

# Long ver.
import math
N, K, p = map(int, input().split())
listN, listK, s=[], [], 1
while N!=0:
    listN.append(N % p)
    listK.append(K % p)
    N, K = N//p, K//p
for i in range(len(ListN)):
    s *= math.comb(ListN[i],ListK[i])
print(s % p)

# Short ver.
import math
N,K,p=map(int,input().split());s=1
while N+K:s*=math.comb(N%p,K%p);N//=p;K//=p
print(s%p)

 

처음으로 숏코딩이라는 걸 하게 된 문제이기도 하다.

Short ver. (101B)는 현재까지 BOJ에 제출된 모든 파이썬계 코드를 통틀어서(Python 2, Python 3, PyPy2, PyPy3) 가장 짧은 코드로 등록되어 있다.

아쉽게도 등록 시점 때문에 1위가 아니라 2위를 하기는 했지만.

 

각설하고, 이 문제를 풀기 위해서는 정수론 지식이 적지 않게 필요하다.

수학적 공대생이라고 대문에서 소개했다시피 내가 수학에 대한 호기심이 적지 않은데, 그렇게 알게 된 뤼카의 정리가 없었으면 이 문제에 손을 댈 생각도 하지 않았을 것이다.

 

우선 조건을 보자. N이 100경 이하의 정수이고, K도 N과 크게 다르지 않은 크기로 입력이 된다.

아무리 M이 작고 Python이 정수의 길이에 신경쓰지 않는 언어이긴 하지만, 문제에 주어진 시간 제한 5초로는 턱도 없다.

이걸 전문 용어로는 시간복잡도가 너무 크다고 하는데, 그에 대하여 설명할 지식은 없으니 패스.

 

시간복잡도를 크게 줄일 수 있는 방법이 상술한 뤼카의 정리이다.

이항계수를 확장해서, n < r일 때 \( \binom{n}{r}=0 \)이라고 정의하자.

그러면 소수 p에 대해서 $m, n=\sum_{i=0}^km_ip^i, \sum_{i=0}^kn_ip^i$로 주어질 시 다음이 성립한다.

\[\binom{m}{n}\equiv\prod_{i=0}^k \binom{m_i}{n_i} \pmod{p}\]

 

우선 문제에서 입력된 자료 N, K, p(문제의 M에 대응함)를 정수에 대응하고 결과가 되어 줄 s, 그리고 $n_i, k_i$를 각각 저장해 줄 listN, listK를 지정하여 주자.

그 다음 while문 안으로 들어가는데, N이 0이 아닐 동안 반복문 내의 문구를 실행하라고 되어 있다. 실행하는 문구는 N, K를 각각 p로 나눈 나머지를 listN, listK의 마지막 원소로 추가하고, N, K를 각각 p로 나눈 몫을 새로운 N, K로 두자는 것이다.

말로만 해서는 이해가 잘 안 될 것이다. 그러면 실제 반복문을 돌리기 전, 1회 돌린 후, 2회 돌린 후의 N, K, listN, listK를 직접 보자.

$$N=\sum_{i=0}^mn_ip^i, K=\sum_{i=0}^mk_ip^i, \text{listN}=[], \text{listK}=[]$$

$$\Downarrow$$

$$N=\sum_{i=1}^mn_ip^{i-1}, K=\sum_{i=1}^mk_ip^{i-1}, \text{listN}=[n_0], \text{listK}=[k_0]$$

$$\Downarrow$$

$$N=\sum_{i=2}^mn_ip^{i-2}, K=\sum_{i=2}^mk_ip^{i-2}, \text{listN}=[n_0,n_1], \text{listK}=[k_0,k_1]$$

 

이제 어느 정도 감이 잡히는가?

N이 while문 안에서 돌며 0이 되어 while문이 종료되는 시점에서, $\text{listN}=[n_0, n_1, \cdots, n_m], \text{listK}=[k_0, k_1, \cdots, k_m]$이 저장되게 된다!

뤼카의 정리를 사용하기 딱 좋게 가공이 되었다는 거다.

남은 것은 아래쪽 for문을 돌면서, s에 $\binom{n_i}{k_i}$를 곱해 준 다음, s를 p로 나눈 나머지를 출력하는 것 뿐이다.

 

뤼카의 정리를 알아도 그것을 코드로 구현하는 것은 쉽지가 않다. 괜히 이 문제가 플래티넘 V를 먹은 것이 아니라는 것이다.

그러나 자기 손으로 한 번 이뤄 보면, 이전까지와는 제법 달라진 본인의 구현 능력을 실감할 수 있을 것이다.

내가 그랬으니까.

 

이로서 11402번의 풀이를 마친다.

 

그럼, 오늘도 당신의 코딩 실력이 상승하기를.