문제 이해하기
Ajob 언어라는 아주 특이한 언어가 있다고 가정해 봅시다. 이 언어는 무한히 많은 글자를 가지고 있으며, 우리는 이 언어로 작성된 n개의 단어를 알고 있습니다.
여기서 흥미로운 규칙이 존재합니다.
- 첫 번째 단어의 길이는 1, 두 번째 단어의 길이는 2와 같이 i번째 단어의 길이는 i입니다.
- 각 단어를 이루는 모든 글자는 서로 중복되지 않습니다(유니크).
이때 n개의 단어 중 하나를 골라 부분 수열(subsequence)을 만들려고 합니다. 조건은 다음과 같습니다.
- 선택한 단어의 길이를 L이라 할 때, 부분 수열의 길이는 반드시 (L − k)여야 합니다.
- 단어의 길이가 k보다 작다면 그 단어는 선택할 수 없습니다.
두 부분 수열은 길이가 다르거나, 같은 위치에 서로 다른 문자가 들어 있을 때 서로 다른 것으로 간주합니다. 우리가 구해야 할 값은 가능한 모든 경우의 수를 소수 p로 나눈 나머지입니다.
예를 들어 입력이 n = 6, k = 5, p = 11이라면 출력은 7이 됩니다.
접근 방법
이 문제의 핵심은 조합론에 있습니다. 길이가 L인 단어에서 (L − k)개의 문자를 남기는 부분 수열을 만드는 방법의 수는 결국 "L개 중 제외할 k개를 고르는 경우의 수", 즉 이항계수 C(L, k)와 같습니다. 따라서 정답은 다음과 같이 정리됩니다.
C(k, k) + C(k+1, k) + … + C(n, k)
그리고 잘 알려진 하키 스틱 항등식(hockey-stick identity)에 의해 이 합은 C(n+1, k+1) 하나로 압축됩니다. 코드에서 n과 k를 각각 1씩 증가시키는 이유가 바로 이것입니다.
남은 과제는 C(n+1, k+1)을 소수 p로 나눈 나머지를 효율적으로 구하는 것입니다. 숫자가 매우 커질 수 있으므로 뤼카 정리(Lucas' Theorem)를 사용합니다. 뤼카 정리는 이항계수를 p진법 자릿수별로 분해해 계산할 수 있게 해주며, 여기에 모듈러 역원과 팩토리얼 전처리를 함께 사용하면 빠른 시간 안에 답을 얻을 수 있습니다.
풀이 절차
- 결과를 저장할 빈 딕셔너리 memo를 생성합니다.
- n := n + 1, k := k + 1로 보정합니다.
- fact := [1]로 시작하는 팩토리얼 리스트를 만듭니다.
- i를 1부터 p−1까지 순회하며 fact의 마지막 원소 × i mod p를 추가합니다.
- p가 memo에 이미 있다면 inv_fact := memo[p]를 사용합니다.
- 그렇지 않다면 다음을 수행합니다.
- inv := [0, 1]로 시작하는 모듈러 역원 리스트를 만듭니다.
- i를 2부터 p−1까지 순회하며 (p − (p // i) × inv[p mod i] mod p)를 추가합니다.
- inv_fact := [1]로 시작해 i를 1부터 p−1까지 순회하며 마지막 원소 × inv[i] mod p를 추가합니다.
- memo[p] := inv_fact로 메모이제이션합니다.
- ret := 1로 초기화합니다.
- n > 0인 동안 다음을 반복합니다.
- n1 := n mod p, k1 := k mod p를 구합니다.
- k1 > n1이라면 해당 자릿수에서 조합이 성립하지 않으므로 0을 반환합니다.
- ret := ret × fact[n1] × inv_fact[k1] × inv_fact[n1 − k1] mod p로 갱신합니다.
- n := n // p, k := k // p로 다음 자릿수로 넘어갑니다.
- ret을 반환합니다.
구현 예시
아래 파이썬 코드를 통해 더 잘 이해해 봅시다.
memo = {}
def solve(n, k, p):
n += 1 # 하키 스틱 항등식 적용을 위한 보정
k += 1
# 팩토리얼 테이블 전처리 (mod p)
fact = [1]
for i in range(1, p):
fact.append(fact[-1] * i % p)
# 역팩토리얼 테이블 (메모이제이션 활용)
if p in memo:
inv_fact = memo[p]
else:
inv = [0, 1] # 모듈러 역원 테이블
for i in range(2, p):
inv.append(p - p // i * inv[p % i] % p)
inv_fact = [1]
for i in range(1, p):
inv_fact.append(inv_fact[-1] * inv[i] % p)
memo[p] = inv_fact
# 뤼카 정리로 이항계수 C(n, k) mod p 계산
ret = 1
while n > 0:
n1 = n % p
k1 = k % p
if k1 > n1: # 어떤 자릿수에서도 k > n이면 조합은 0
return 0
ret = ret * fact[n1] * inv_fact[k1] * inv_fact[n1 - k1] % p
n //= p
k //= p
return ret
n = 6
k = 5
p = 11
print(solve(n, k, p))입력
6, 5, 11
출력
7
마무리
이 풀이의 시간 복잡도는 전처리에 O(p), 뤼카 정리 적용에 O(logp n)이므로 n이 매우 커져도 안정적인 성능을 보장합니다. 조합 문제에서 답을 큰 소수로 나눈 나머지를 요구한다면, 뤼카 정리와 모듈러 역원 전처리를 조합하는 이 패턴을 기억해 두면 다양한 문제에 응용할 수 있습니다.