September 23, 2023 00:00

문제 상황

아래의 문제 상황은 BOJ 16216번 문제에서 만나볼 수 있다.

정점이 가중치가 없는 트리 $T$가 주어진다. $T$는 $N$개의 정점과 $N-1$개의 간선을 가지고 있다. 정점들 가운데, 출발 정점 $v$와, 방문하고 싶은 $K$개의 정점들 $u_1, \cdots, u_K$이 주어진다. 정점 $v$에서 출발해서, $u_1, \cdots, u_K$ 중 1개, 2개, …, K개를 방문하는 가장 짧은 경로의 길이를 각각 구하여라.

아래부터 이 문제에 대한 풀이로 개념에 대한 설명을 시작하려 하니, 아직 이 문제에 대해 고민해보지 않았다면 (스포일러를 피하고 싶다면) 스크롤을 내리지 않기를 권한다. 아래 단락에 서술할 풀이 자체도 재미있다.

$O(NK)$ 발상

이 문제에 흔히 알려진 “트리 압축” 기법을 사용하기 위해서는 아래의 사고 과정을 거쳐와야 한다. 아래의 풀이로 이 문제를 $O(NK)$에 풀 수 있다.

• 정점 $v$를 트리의 루트로 두고 생각하자.
• 문제에서 찾는 경로 ($K$개의 정점들 중 $i$ ($1 \le i \le K$)개를 방문하는 경로)는 정점 $v$에서 시작하는 DFS와 비슷한 형태가 될 것이다.
• 구체적으로, 어떤 서브트리에 대해서, 이 서브트리를 방문하는 횟수는 최대 1번이다. 어떤 서브트리에 들어갔다가 나오는 행위를 두 번 이상 반복하지 않는다는 의미이다.
• 어떤 정점 $x$의 자녀 정점 $c_1, \cdots, c_{\mathrm{childs}(x)}$에 대해, 각 $c_i$를 루트로 하는 서브트리에 대한 문제의 답은 모두 독립이다. 즉 트리 DP로 문제를 풀 수 있다.

이제 다음과 같이 DP 배열과 점화식을 정의할 수 있다. 마지막에 출발 정점으로 돌아올 필요가 없다는 조건 때문에, 두 개의 DP 배열을 정의해야 한다.

• $\mathrm{dp}_1(x, i)$ := 정점 $x$에서 출발해서, $x$를 루트로 하는 서브트리 안에서 찾고 싶은 정점 $i$개를 방문하고 $x$로 돌아오는 최단 경로의 길이

• $\mathrm{dp}_2(x, i)$ := 정점 $x$에서 출발해서, $x$를 루트로 하는 서브트리 안에서 찾고 싶은 정점 $i$개를 방문하는 ($x$로 돌아올 필요는 없음!) 최단 경로의 길이

$\mathrm{dp}(c_k, j_k)$ ($1 \le k \le \mathrm{childs}(x)$)의 값들을 이용해 $\mathrm{dp}(x, i)$의 값을 구할 수 있다. $j_1 + \cdots + j_{\mathrm{childs}(x)} = i$일 때에, 각 자식 정점을 루트로 하는 서브트리에서 만들어진 경로들을 합쳐서 총 $i$개의 ‘방문하고 싶은 정점’을 방문하는 하나의 경로를 만들 수 있게 된다. 자식 정점을 하나씩 방문하면서, 이들의 $\mathrm{dp}$ 값들을 $x$의 $\mathrm{dp}$ 배열에 업데이트해주는 방법으로 진행하게 된다. $k$번째 자식 정점 $c_k$를 방문했을 때, 다음과 같은 점화식을 통해 $\mathrm{dp}$ 배열을 업데이트하게 된다. (아래 식에는 $x$가 ‘방문하고 싶은 정점’일 때의 처리와 여러 가지 예외 케이스들에 대한 처리가 생략되어 있다)

• $\mathrm{dp}1 (x, i+j) \Leftarrow \mathrm{dp}_1(x, i) + \mathrm{dp}_1(c{k}(x), j) + \cdots$

• $\mathrm{dp}_2 (x, i+j) \Leftarrow \min(\mathrm{dp}_1(x, i) + \mathrm{dp}_2(c_k(x), j), \mathrm{dp}_2(x, i) + \mathrm{dp}_1(c_k(x), j)) + \cdots$

이와 같이 각 정점 $x$의 자식 정점들을 순회하면서 $\mathrm{dp}(x, i)$ ($0 \le i$, $i$는 $x$를 루트로 하는 서브트리의 ‘방문하고 싶은 정점’의 개수 이하) 들을 모두 채우는 데 드는 시간은 모두 합쳐서 $O(NK)$가 된다. 각 정점을 방문할 때마다 2중 반복문을 돌려야 하므로 그렇게 되지 않을 것 같지만 놀랍게도 그건 사실이다. 이에 대한 증명도 재미있으니, 대략적으로 한 번 수행해보기 바란다.

트리 압축

위 풀이는 DP 풀이이므로, 본질적으로 모든 정점에 대해 $\mathrm{dp}$ 배열 값을 구해야 한다. 채워야 하는 값의 개수가 최대 $O(NK)$개이므로, $O(NK)$ 미만의 시간에 문제를 푸는 것은 일견 불가능해 보인다. 이를 개선하는 데에 트리 압축을 사용할 수 있다.

우리는 전체 $N$개의 정점들 중 $K$개의 정점에만 관심이 있다. 나머지 정점들은 트리에서 ‘생략’해도 된다. 간선들을 합치는 방법으로, 트리에서 $K$개의 정점의 위상을 그대로 유지하면서 $O(K)$개의 정점만을 남길 수 있으며, 이 기법을 ‘트리 압축’ 또는 virtual tree, auxiliary tree 등의 이름으로 부른다. 다음 두 조건 중 하나를 만족하는 정점을 남긴다.

• 자기 자신이 관심 있는 정점이다.
• 자식 정점들 중 서로 다른 2개 이상이 다음을 만족한다: 이 정점을 루트로 하는 서브트리 안에 관심 있는 정점이 포함되어 있다.

트리의 오일러 회로(DFS와 동일하다)를 응용한 방법으로 이를 수행하는 방법은 이 게시글에 자세히 설명되어 있다.

‘방문하고 싶은 $K$개의 정점’들에 대해 트리 압축을 수행하면 입력 트리의 정점의 개수가 $O(K)$로 줄어들며, 이제 위에 서술한 알고리즘을 $O(K^2)$에 수행할 수 있게 된다. 이 방법으로 시간 안에 문제의 답을 구할 수 있게 된다.

하지만, 사실 이와 같은 개념을 이해하고 있기만 하다면, 트리 압축을 명시적으로 수행할 필요가 없음을 알게 된다. 원래의 트리에서 동일한 DP를 수행한다고 하자. 압축된 트리에서 사라질 정점 (위 조건을 만족하지 않는 정점)의 DP 배열은, 자식 정점들 중 하나의 DP 배열과 같다. 따라서, 이 때는 해당 배열의 레퍼런스(포인터)를 복사하는 방법으로, 추가적인 시간을 들이지 않고 해당 정점의 DP 배열을 채울 수 있다. 트리 압축 전처리와 본질적으로는 동일한 방법이지만, 자주 등장하는 small-to-large 최적화와 유사한 방법으로 같은 효과를 낼 수 있다. 시간 복잡도는 $O(N+K^2)$이다.

유사한 문제

BOJ 20535는 트리 압축을 이용해 풀 수 있는 유명한 연습문제이다. 이 문제에서는 ‘관심 있는’ 정점이 쿼리에 따라 바뀌고, 심지어 쿼리들의 ‘관심 있는’ 정점의 합집합의 크기는 최대 100만이므로, 트리 압축을 전처리로 수행할 수 없다. 이 문제에서도 트리 압축을 동적으로 또는 암묵적으로 수행해야 하는데, 여기서는 트리의 오일러 회로를 이용한 트리 압축의 정의가 빛을 발한다.

쿼리로 $K$개의 정점 $q_1, \cdots, q_K$가 주어졌다고 하자. 이 쿼리에서 고려해야 하는 트리는, 이들 $K$개의 정점들과, $q_i$와 $q_j$의 LCA들만 남긴 트리이다. $K$개의 정점을 ‘오일러 투어 순서’ (원래 트리의 중위 순회 순서)로 정렬하고 나면, $q_1$과 $q_K$의 LCA가 전체 정점들의 LCA와 같음을 알 수 있다. 전체 문제의 답을 $f(1, K)$ (i.e. $q_{1..K}$에 대한 답)이라고 두면, 전체의 LCA ($LCA(q_1, q_K)$)가 답이 되지 않는 쌍의 개수를 재귀적으로 구해 내려가는 방법으로 $O(K)$에 쿼리에 대한 답을 구할 수 있다.