문제 소개
n개의 촛불이 왼쪽에서 오른쪽으로 일렬로 배치되어 있다고 가정해 봅시다. 왼쪽에서 i번째 촛불은 높이 h[i]와 색상 c[i]를 가지며, 정수 k는 1부터 k까지의 색상이 존재함을 나타냅니다. 우리가 구해야 하는 것은 '엄격하게 증가하는 다채로운(colorful) 시퀀스'의 총 개수입니다.
여기서 증가 여부는 촛불의 높이를 기준으로 판단하며, 1부터 k까지의 모든 색상에 해당하는 촛불이 각각 최소 하나씩 포함되어 있을 때 그 시퀀스를 '다채로운(colorful)' 시퀀스라고 정의합니다. 답이 매우 커질 수 있으므로 최종 결과는 10^9 + 7(= 1000000007)로 나눈 나머지를 반환해야 합니다.
예시
예를 들어 K = 3, h = [1, 3, 2, 4], c = [1, 2, 2, 3]이 주어진다면 출력은 2입니다. 조건을 만족하는 시퀀스는 [1, 2, 4]와 [1, 3, 4], 단 두 개뿐이기 때문입니다.
접근 방법: 펜윅 트리와 포함-배제의 원리
이 문제는 펜윅 트리(Fenwick Tree, Binary Indexed Tree)로 증가 부분 수열의 개수를 세고, 비트마스크를 활용한 포함-배제(inclusion-exclusion) 기법으로 모든 색상을 포함하는 경우만 남기는 방식으로 해결할 수 있습니다. 핵심 단계는 다음과 같습니다.
- read(T, i) : 펜윅 트리 T에서 인덱스 i까지의 구간합을 구하는 함수입니다. i에서 최하위 비트(i AND -i)를 반복적으로 빼 가면서 값을 누적하고, 매번 10^9 + 7로 나눈 나머지를 유지한 뒤 그 합을 반환합니다.
- update(T, i, v) : 펜윅 트리의 i번째 위치에 값 v를 반영하는 함수입니다. i에 최하위 비트를 계속 더해 가며 상위 노드들을 갱신하고, 역시 10^9 + 7로 나눈 나머지를 저장합니다.
- 메인 로직 : L := 2^k, R := 0, N := len(h)로 초기화한 뒤, 0부터 L-1까지의 각 비트마스크 i에 대해 아래 과정을 수행합니다.
- 크기 50010의 배열 T를 0으로 초기화하고, 누적 변수 t := 0으로 설정합니다.
- j = 0부터 N-1까지 순회하면서 (i >> (c[j] - 1)) & 1이 참, 즉 현재 마스크에 j번째 촛불의 색상이 포함되어 있다면 update(T, h[j], read(T, h[j] - 1) + 1)을 호출해 h[j]로 끝나는 증가 시퀀스의 개수를 갱신하고 그 값을 t에 더합니다.
- 마스크 i에 설정된 비트 개수의 홀짝성이 k의 홀짝성과 같으면 R에 t를 더하고, 그렇지 않으면 R에서 t를 뺍니다(항상 10^9 + 7 기준의 모듈로 연산을 유지). 이것이 포함-배제 원리의 핵심입니다.
- 모든 마스크를 처리한 뒤 R을 반환하면, 그것이 곧 조건을 만족하는 시퀀스의 총 개수입니다.
이 방식의 시간 복잡도는 O(2^k × N log H)로, 색상의 종류 k가 작을 때 매우 효율적으로 동작합니다.
파이썬 구현 코드
아래 구현을 직접 실행해 보면 동작 방식을 더 잘 이해할 수 있습니다.
def solve(k, h, c):
def read(T, i):
s = 0
while i > 0:
s += T[i]
s %= 1000000007
i -= (i & -i)
return s
def update(T, i, v):
while i <= 50010:
T[i] += v
T[i] %= 1000000007
i += (i & -i)
return v
def number_of_bits(b):
cnt = 0
while b:
b &= b - 1
cnt += 1
return cnt
L = 2 ** k
R = 0
N = len(h)
for i in range(L):
T = [0 for _ in range(50010)]
t = 0
for j in range(N):
if (i >> (c[j] - 1)) & 1:
t += update(T, h[j], read(T, h[j] - 1) + 1)
t %= 1000000007
if number_of_bits(i) % 2 == k % 2:
R += t
R %= 1000000007
else:
R += 1000000007 - t
R %= 1000000007
return R
k = 3
h = [1, 3, 2, 4]
c = [1, 2, 2, 3]
print(solve(k, h, c))
실행 결과 확인
위 코드에 아래 입력을 넣어 실행하면:
3, [1, 3, 2, 4], [1, 2, 2, 3]
다음과 같은 출력을 얻을 수 있습니다.
2
k = 3일 때 가능한 모든 색상 조합(비트마스크)에 대해 증가 시퀀스를 세고 포함-배제로 정확히 걸러낸 결과, 조건을 만족하는 시퀀스는 [1, 2, 4]와 [1, 3, 4] 두 개뿐임을 확인할 수 있습니다.