백준 12969 - ABC

문제 설명

입출력 및 제한 사항

풀이

별다른 설정이 없이 굉장히 간단명료한 문제 설명이다. 하지만 풀이 과정은 간단하지 않다...

요약해보자면 길이 n의 문자열 s를 {A, B, C}로만 만들 때, 오름차순 쌍의 총 개수가 정확히 k가 되도록 하는 임의의 s 하나를 출력하는 문제이다. 없다면 -1을 출력한다.

* 오름차순 쌍: 인덱스 i < j이면서 S[i] < S[j] (사전순 기준 A < B < C)

 

예를 들면 다음과 같다.

"ABC"라면 쌍은 (A, B), (A, C), (B, C) 로 총 3개

"CBA"라면 오름차순 쌍은 0개

 

문제를 관찰해보자

문자 하나를 왼쪽부터 차례로 붙인다고 할 때, 지금까지 만든 부분 문자열에 들어 있는 A와 B의 개수만 알면, 다음 문자를 붙일 때 쌍이 얼마나 증가하는지 결정된다.

• A를 붙이면 새로운 쌍 없음
• B를 붙이면 앞에 있던 A의 개수만큼 새로운 쌍 생김
• C를 붙이면 앞에 있던 A와 B의 개수만큼 새로운 쌍 생김

이 문제는 같은 상태가 여러 경로로 반복해서 등장하는 중복 부분 문제가 있고, 그 상태에서 “해가 존재하는지/어떤 접미 문자열이 되는지”가 이후 선택에만 의존하는 부분 구조로 DAG를 성립하여 DP문제라고 볼 수 있다. 해당 조건들을 활용하여 구현해보자. 

내 코드

import sys
input = sys.stdin.readline
mis = lambda: map(int, input().split())

n, k = map(int, input().split())
MASK = (1 << (k+1)) - 1

dp = [[[0] * (n+1) for _ in range(n+1)] for _ in range(n+1)]
for a in range(n+1):
    for b in range(n+1):
        dp[0][a][b] = 1

for rem in range(1, n+1):
    for a in range(n+1):
        for b in range(n+1):
            bits = 0
            
            if a + 1 <= n:
                bits |= dp[rem-1][a+1][b]
                
            if b + 1 <= n:
                bits |= (dp[rem-1][a][b+1] << a)
            
            bits |= (dp[rem-1][a][b] << (a+b))
            dp[rem][a][b] = bits & MASK

need = k
a = b = 0
ans = ""
for i in range(n):
    rem = n - i - 1

    if a + 1 <= n and ((dp[rem][a+1][b] >> need) & 1):
        ans += 'A'
        a += 1
        continue

    need2 = need - a
    if b + 1 <= n and need2 >= 0 and ((dp[rem][a][b+1] >> need2) & 1):
        ans += 'B'
        b += 1
        need = need2
        continue

    need2 = need - (a + b)
    if need2 >= 0 and ((dp[rem][a][b] >> need2) & 1):
        ans += 'C'
        need = need2
        continue

    print(-1)
    exit()

print(ans)

 

dp 테이블은 다음과 같이 정의했다.

dp[rem][a][b]: 현재 A, B 개수에서 앞으로 rem칸을 채울 때 만들 수 있는 추가 쌍들의 비트셋

비트셋에서 정수의 t번째 비트가 1이면 추가 쌍을 정확히 t개 더 만드는 게 가능하다는 것을 표기한다.

 

우리는 a, b, c를 추가했을 때 개수가 어떻게 늘어나는지 알고 있으므로 쉽게 확인가능하다. 여기서 비트셋을 사용한 이유는 일반 정수로는 가능한 추가 쌍 수의 모든 경우를 표현할 수 없기 때문이다.

예를 들어 rem=3, a=2, b=0이면 실제로 가능한 추가 쌍 수 집합이 {0, 2, 3, 4, 5, 6, 7, 8} 처럼 비연속적일 수 있다.

 

위 dp 테이블을 순회하며 각 상태에서 a, b, c가 증가한 경우 모두 추가해주고 MASK와 AND연산해주어 넘어가는 집합들은 제거해주었다. 여기서 MASK는 모든 가능한 상태쌍의 집합이다.

 

이후 문자열 구성 단계에선 매 위치 i(0부터 시작)마다 세가지 값만 관리한다.

• a : 지금까지 넣은 A의 개수
• b : 지금까지 넣은 B의 개수
• need : 앞으로 더 만들어야 할 오름차순 쌍 개수 (초기값 k)

지금 자리에서 문자를 하나 두면 즉시 생기는 쌍 증가량은 위에서 설명한 것과 같고 “이번에 X만큼 생성가능하니, 나머지 rem칸으로 정확히 need-X를 만들 수 있는가?”를 dp로 확인한다.

A를 둘 수 있나?

X=0이라 need는 그대로 유지

가능 조건: (F[rem][a+1][b]에 need 비트가 켜져 있음) 즉, ((F[rem][a+1][b] >> need) & 1) == 1

B를 둘 수 있나?

X = a → 남은 필요 need2 = need - a

need = need2로 갱신

가능 조건: need2 >= 0 그리고 ((F[rem][a][b+1] >> need2) & 1) == 1

C를 둘 수 있나?

C는 A, B 개수에 변화가 없으니 (a, b) 그대로

Δ = a + b → need2 = need - (a + b)

need = need2로 갱신

가능 조건: need2 >= 0 그리고 ((F[rem][a][b] >> need2) & 1) == 1

 

셋 다 불가능하면 어떤 선택을 해도 목표를 달성 못 한다는 뜻이므로 -1을 출력하고 종료한다.

 

시간 복잡도는 dp테이블 구성에 O(n^3), 문자열 생성에 O(n)이므로 O(n^3)이라고 볼 수 있다.

고찰

DFS를 사용한 탑다운 DP 방식으로도 풀이할 수 있다.

from functools import lru_cache
import sys
input = sys.stdin.readline
mis = lambda: map(int, input().split())
sys.setrecursionlimit(10**5)

n, k = mis()

@lru_cache(maxsize= None)
def dfs(pos, a, b, cnt):
    if cnt > k:
        return None
    
    r = n - pos
    max_add = r * (a + b) + r * (r - 1) // 2
    if cnt + max_add < k:
        return None

    if pos == n:
        return "" if cnt == k else None

    t = dfs(pos+1, a+1, b, cnt)
    if t is not None:
        return "A" + t

    t = dfs(pos+1, a, b+1, cnt+a)
    if t is not None:
        return "B" + t

    t = dfs(pos+1, a, b, cnt+a+b)
    if t is not None:
        return "C" + t

    return None

ans = dfs(0, 0, 0, 0)
print(ans if ans is not None else -1)

각 변수들의 의미는 다음과 같다.

• pos: 지금까지 만든 길이 (다음에 넣을 위치)
• a: 지금까지 넣은 A의 개수
• b: 지금까지 넣은 B의 개수 (C는 pos - a - b)
• cnt: 지금까지 누적된 오름차순 쌍의 개수

또한 두 가지의 가지치기 조건으로 효율을 올려주었다.

1. 초과 컷

쌍의 개수는 감소하지 않기 때문에 cnt > k이면 더 이상 진행해도 k로 내려올 방법이 없다. 따라서 해당 분기에서 쌍의 개수가 k개를 넘는다면 이후의 분기는 탐색하지 않아도 된다.

 

2. 상한 컷 (최대로 벌어도 k 미만이면 중단)

현재 문자열에서 쌍이 증가할 수 있는 최대 문자열은 남은 자리들에 전부 C만 붙이는 것이다.

해당 경우에 첫 글자는 a+b만큼, 다음은 a+b+1, …, 마지막은 a+b+(r-1) 만큼 증가 가능하다.

따라서 남은 칸은 r = n - pos이고, 앞으로 최대로 벌 수 있는 추가 쌍의 상한은 r(a+b) + (r(r−1) / 2​​)이다. 따라서 해당 상한 값을 계산한 결과가 k보다 작다면 이 분기는 더 이상 탐색하지 않아도 된다.

 

시간 복잡도를 계산해보면 a를 선택하거나, b를 선택하거나, c를 선택하거나 이 세 가지 분기로 나누어지기 때문에 O(n^3)이라고 착각할 수 있다. 하지만 실제 계산은 각 상태에서 이루어지기 때문에 다르다. 작성한 코드는 상태를 (pos, a, b, cnt)로 두고 pos는 이전까지 만든 단어의 길이이기 때문에 중복된 계산이 나올 수 있다.

 

정확한 계산은 다음과 같다.

1. 고정된 길이 pos에서

a ≥ 0, b ≥ 0, a + b ≤ pos

를 만족하는 정수 쌍 (a,b)의 개수를 센다. (c는 굳이 세지 않아도 c = pos − a − b로 결정되기 때문에 a와 b의 개수만 센다.)

-> O(n^2)

2. pos=0..n을 합치면 O(n^3)

3. DP 상태에 cnt(0..K)까지 포함하면 O(n^3 * K)로 확장

-> 분기 A, B, C는 각 상태당 상수 작업이므로 차수에 영향 없음

따라서 정확한 시간복잡도는 O(n^3 * K)이다. 하지만 가지치기로 인해 실제 연산은 많이 줄어들어 바텀업 DP보다 적은 시간에 연산이 완료된 것을 확인할 수 있다.