BJ 11051 이항 계수 2

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

이항 계수의 성질

$n$개의 원소를 가지는 집합에서 $k$개를 택하는 경우의 수는 이들 중 원소의 개수가 $(n-1)$인 부분 집합 하나를 택하여 여기서 $(k-1)$개를 택하고 남은 원 소 하나를 택하는 경우의 수와 원소의 개수가 $(n-1)$인 부분 집합에서 전체 $k$개를 택하는 경우의 수의 합과 같다는 사실을 알 수 있습니다.

DP와 memoization

위 성질을 이용하여 이항 계수를 구하는 함수를 작성해 봅시다.

def choose(n, k):
	if n == 1 or k == 0 or k == n:
    	return 1
    return choose(n-1, k-1) + choose(n-1, k)

이 경우 재귀 함수의 성질에 따라 choose 함수는 두 개의 자식 함수를 호출합니다.

재귀의 깊이가 n 까지 임을 생각하면 총 호출해야 하는 함수는 $1+2+4+8+\cdots+2^n=2^{n+1}-1$ 개이므로, 주어진 시간 내에 해결이 불가능할 것입니다.

이들 중 겹치는 부분이 존재하는데, 이를 계산해 놓고 바로 사용할 수 있게 하면 계산량을 줄일 수 있습니다. 이를 memoization이라고 하며, 계산량을 $O(n^2)$ 수준으로 줄일 수 있습니다.

코드

import sys
sys.setrecursionlimit(2000)

memo = [[-1] * 1001 for _ in range(1001)]

def choose(n, k):
    if n == 1 or k == 0 or k == n:
        return 1

    if memo[n][k] == -1:
        memo[n][k] = (choose(n-1, k-1) + choose(n-1, k)) % 10007
    return memo[n][k]

n, k = map(int, input().split())
print(choose(n, k))

구현시 주의할 점

배열의 인덱싱

문제에서 제시하는 최대 범위인 $1000 \choose k$에 대해, 본 코드는 memo[1000][1000]을 접근합니다. 따라서 memo 배열, 정확히는 인덱스로 접근할 수 있는 리스트의 원소 개수는 1001이어야 합니다(인덱싱은 0부터 시작하니까요).

따라서 memo를 정의할 땐 1001×1001 크기의 배열을 만들어야 합니다.

또한 이와 같은 한 줄의 정의를 list comprehension이라고 합니다.

재귀 깊이

Python의 기본 재귀 깊이 최대치는 1000입니다. 이는 PyPy를 통해 완화할 수 있으나, PyPy는 재귀에 상당히 약하다는 단점이 있습니다.
이는 sys 라이브러리의 setrecursionlimit 메서드를 이용해 재귀 한계 깊이를 설정할 수 있습니다.

본 코드에서는 넉넉하게 2000으로 설정하였습니다.