ABC340 G - Leaf Color
Competitive Programming C++ AtCoder Algorithms
2024-02-14
공식 에디토리얼
문제 링크
https://atcoder.jp/contests/abc340/tasks/abc340_g
원문 링크
https://atcoder.jp/contests/abc340/editorial/9249
먼저 이 문제는 몇 가지의 접근법이 있지만, 이 해설에서는 Auxiliary Tree 라는 테크닉을 소개합니다.
접근법
먼저 이 문제의 접근법에 대해 설명하겠습니다.
차수가 \(1\)인 정점에 칠해진 색 \(c\)를 고정합니다. 그렇게 하면 고정했을 때의 문제는 Tree DP로 \(O(N)\)에 풀 수 있습니다. (이 방법은 쉽지만은 않으므로 Blue diff 급 난이도 정도입니다) 하지만 모든 색에 대해 계산하자면 \(O(N^2)\)가 걸려 시간 초과를 받게 됩니다.
색 \(c\)를 고정했을 때의 문제가 \(O(N)\)이 걸리는 이유는 트리의 정점이 \(N\)개 있기 때문입니다. 하지만 \(N\)개의 정점들 중에는 절대로 사용되지 않는 정점이나 원래부터 필요 없는 정점이 많이 있어, 색 \(c\)에 대해 DP를 할 때 이런 정점들을 포함해서 계산하게 된다면 낭비가 일어나게 됩니다. 여기서 트리를 압축하는 것을 생각할 수 있습니다.
트리를 압축할 때 남길 정점들을 생각해봅시다. 색 \(c\)인 정점은 무조건 남긴다고 치고, 정점끼리의 위치 관계를 보전하기 위해서는 색 \(c\)인 정점끼리의 \(\mathrm{LCA}\)가 되는 정점도 남겨야 할 듯 합니다. (여기서 트리 \(T\)는 정점 \(1\)을 루트로 하는 트리로 생각합니다.) 이런 식으로 필요한 정점들을 골라 자손 관계인 부분에 간선을 연결해주면 왼쪽 그림에서 오른쪽 그림처럼 됩니다. 오른쪽 그림의 상태가 된 트리에서 DP를 해도 문제 없을 듯 합니다.
(그림: 정점 \(6, 7, 8, 14, 15\)가 색 \(c\)로 칠해졌을 때의 압축 방법)
이런 식으로 트리를 압축해서 만들어지는 트리를 지정된 정점들의 최소공통조상을 유지한 채로 트리를 압축해서 만들 수 있는 보조적인 트리라고 부릅니다. (일반적으로 Auxiliary Tree, Virtual tree, 트리 압축이라고 불리는 경우도 있습니다. 이하 Auxiliary Tree라 하겠습니다.)
이름에 유래에 대해서는 다음을 참고해 주세요.
Auxiliarty Tree에 대한 Kyopro Encyclopedia of Algorithms의 글
Auxiliarty Tree 의 성질
Auxiliarty Tree를 더 형식적으로 설명하면 다음과 같습니다. 루트가 있는 트리 \(T\) 및 정점들으 ㅣ집합 \(X\)가 주어졌을 때, 정점들의 집합 \(A(X)\)룰
\[A(X) = \{\mathrm{lca}(x, y) \vert x \in X, y \in X\}\]
라 정의하고, \(A(X)\)에 포함되는 정점들 사이에 \(T\)에서의 자손 관계인 부분의 간선을 추가해서 만들어지는 루트가 있는 트리를 Auxiliarty Tree라고 합니다.
Auxiliarty Tree의 성질로써 다음과 같은 것들을 들 수 있습니다.
성질 1
\(\vert A(X) \vert \leq 2 \vert X \vert - 1\) 이다.
(증명) 트리 \(T\)에 대해 DFS를 한 번 진행합니다. 그 다음 \(X\)의 정점들을 방문 순서대로 순번을 매긴 수열을 \(x_1, x_2, … , x_m\)이라 합니다.
\(l = \mathrm{lca}(x_1, x_2)\)라 하면 다음과 같은 사실이 성립합니다.
보제 1
\(3\)이상의 \(n\)에 대해 \(\mathrm{lca}(x_1, x_n) \neq l\)이라면 \(\mathrm{lca}(x_1, x_n) = \mathrm{lca}(x_2, x_n)\)이 성립한다.
(보제의 증명) 귀류법으로 증명합니다. \(\mathrm{lca}(x_1, x_n) \neq l\)이고 \(\mathrm{lca}(x_1, x_n) \neq \mathrm{lca}(x_2, x_n)\)을 이 만족하는 \(n\)이 존재한다 가정해봅시다. \(\mathrm{lca}(x_1, x_n) = m\)이라 하면 \(m\)은 \(l\)의 자손 또는 조상이 되지만, \(m\)이 \(l\)의 조상이라 하면 \(\mathrm{lca}(x_1, x_n) = \mathrm{lca}(x_2, x_n) = m\)이 되어 모순입니다. 따라서 생각할 필요가 있는 경우는 \(m\)이 \(l\)의 자손일 때지만, 이 때의 방문 순서는 \(x_1, x_n, x_2\)의 순서가 되어 \(x\)의 순서에 대한 조건에 위배됩니다. 따라서 모순입니다.(보제 증명 끝)
보제로부터 \(\mathrm{lca}(x_1, x_n)\)은 (1) \(x_1\)자신인 경우, (2)\(\mathrm{lca}(x_1, x_2)\)인 경우, (3) \(\mathrm{lca}(x_2, x_n)\)과 일치하는 경우 의 3가지 경우밖에 없음을 알 수 있습니다 .따라서 \(A(x)\)는
\[A(\{x_1, …, x_m\}) = A(\{x_2, …, x_m\}) \cup \{x_1, \mathrm{lca}(x_1, x_2)\}\]
로 표현할 수 있는 것을 알 수 있습니다. 이 식으로부터
\[\vert A(X) \vert \leq \vert A(\{x_2, …, x_m\}) \vert + 2\]
가 성립함을 알 수 있으므로, 위의 식을 기초로 정점 수에 대한 귀납법을 돌리면 \(\vert A(X) \vert \leq 2 \vert X \vert - 1\)가 됨을 증명합니다. (증명 끝)
- 참고로 부등식 우변의 \(2 \vert X \vert - 1\)이 제일 엄격한 판정임도 증명할 수 있습니다. 이는 \(\vert A(X) \vert\)가 최대 \(2 \vert X \vert - 1\)개의 정점을 가질 수 있음을 증명하면 되며, iol T가 완전 이진 트리이고 \(X\)가 리프 노드인 정점들로만 구성된 집합임을 생각하면 \(\vert A(X) \vert = 2 \vert X \vert - 1\)인 경우를 생각할 수 있습니다.
이 성질은 Auxiliarty Tree의 정점 수는 \(O(\vert X \vert)\), 측 원본이 되는 정점들의 집합의 크기의 정수 배로 줄어든다는 것을 의미합니다. 원래 문제의 풀이는 이 사실을 이용해 계산량을 줄여나갑니다.
Auxiliarty Tree의 구현
정점들의 집합 \(X\)에 대한 Auxiliarty Tree의 구현은 \(1\)번의 생성마다 \(O(\vert X \vert (\log \vert X \vert + \log N))\)에 할 수 있습니다. (\(N\)은 원본이 되는 루트 있는 트리의 정점의 수. 또, 트리 \(T\)에 대한 전처리에 \(O(N \log N)\)정도 필요합니다.)
구현 방법에 대해 설명하겠습니다. 다음 성질을 이용합니다.
성질 2
트리 \(T\)에 대해 DFS를 한 번 진행한 다음 \(X\)의 정점들을 방문 순서대로 순번을 매긴 수열을 \(x_1, x_2, … , x_m\)이라 한다.
이 때 \(A(X)\)에 대해 다음 사실이 성립한다. \[A(X) = X \cup \{\mathrm{lca}(x_i, x_{i + 1}) \vert 1 \leq i \leq m\}\]
(증명) 성질 1의 증명으로부터 나온
\[A(\{x_1, …, x_m\}) = A(\{x_2, …, x_m\}) \cup \{x_1, \mathrm{lca}(x_1, x_2)\}\]
라는 식을 이용하면
\[A(\{x_1, …, x_m\}) \newline = A(\{x_2, …, x_m\}) \cup \{x_1, \mathrm{lca}(x_1, x_2)\} \newline = A(\{x_3, …, x_m\}) \cup \{x_1, \mathrm{lca}(x_1, x_2), x_2, \mathrm{lca}(x_2, x_3)\} \newline \vdots \newline = A(\{x_m\}) \cup \{x_1, \mathrm{lca}(x_1, x_2), x_2, …, x_{m - 1}, \mathrm{lca}(x_{m - 1}, x_m), x_m\}\]
이 되는 것을 알 수 있습니다.(증명 끝)
성질 2를 이용하면 다음과 같은 순서로 \(A(X)\)에 등장하는 정점들을 방문 순서대로 나열한 수열을 얻을 수 있습니다.
- 전처리 단계에서 HLD 등을 이용해 트리 \(T\)의 \(\mathrm{LCA}\)를 빠르게 구할 수 있도록 한다.
- \(X\)에 등장하는 정점을 방문 순서대로 \(x_1, x_2, …, x_m\)이라 한다.
- \(\mathrm{lca}(x_1, x_2), …, \mathrm{lca}(x_{m - 1}, x_m)\)을 계산한다.
- \(x_1, x_2, …, x_m, \mathrm{lca}(x_1, x_2), …, \mathrm{lca}(x_{m - 1}, x_m)\)를 방문 순서대로 정렬해 중복을 제거한다.
등장하는 정점들을 모두 알아내면 다음은 이 수열로 스택을 이용해서 정점을 이어 나가는 것으로 Auxiliary Tree를 생성할 수 있습니다. (자세히는 yaketake08님의 블로그 글의 「ソート2回の方法」(정렬 2번 하는 방법)을 참고해 주세요.)
이상의 내용으로부터 각 색으로 칠해진 정점의 집합으로부터 Auxiliary Tree를 생성해 DP로 풀면 \(O(N \log N)\)에 이 문제를 풀 수 있습니다.
트리 압축은 무섭다
핵심은 트리 압축과 LCA 구하기인데, HLD까지 필요한 저 구현들은 지금 내 수준에선 시간 내엔 도무지 안 되었을 것이다.
G는 역시 무섭구먼.