APIO 2021
algorithms, , ioi, , data , structure, , geometry
APIO 2021
APIO 2021은 아시아 태평양 학생들을 대상으로 진행되는 정보 올림피아드 대회로, 국제 정보 올림피아드 (IOI) 바로 다음 순위의 권위와 중요성을 가진 국제 대회이다. 조금 늦기는 했지만 아직 APIO 2021의 문제들을 완전히 풀이한 글이 없어서 풀이를 작성한다.
문제 풀이 한 줄 요약
- 육각형 영역: 다각형 내부의 모든 격자점 간의 거리 합을 구하는 문제로 환원할 수 있다. 다각형의 3개의 축을 분리하여 생각하면, 각 축에 대해서 삼각분할과 비슷한 일종의 트리를 형성할 수 있고, 모든 경로가 이 트리 상의 유일한 경로로 환원됨을 관찰할 수 있다. BST에 Plane sweeping을 사용하여 이러한 트리를 빠르게 구성한 후 트리 안에서 동적 계획법을 수행할 수 있다.
- 밀림 점프: Cartesian tree 형태의 그래프에서 최단 경로 쿼리를 빠르게 해결하는 문제이다. 기본적으로 그래프의 Treewidth가 작기 때문에, Separator를 구하는 식으로도 어떻게든 해결할 수 있는 문제이다. 실제로는 그렇게 어렵게 풀 필요는 없고, 그래프의 성질을 관찰한 그리디 알고리즘을 자료구조로 최적화하는 정도로 충분하다.
- 도로 폐쇄: $k = 1$ 일 경우 매칭을 푸는 문제와 동일한데, 이를 일반화하여 “모든 정점 차수가 특정 수 이하인 간선 부분집합” 을 구하는 문제로 바라본 문제이다. $k = 1$ 인 경우의 접근을 일반화하면 쉽게 제곱 시간 풀이를 유도할 수 있다. 이 풀이에서 유용한 전이만을 잘 추려서 업데이트할 경우 총 업데이트의 경우의 수가 $O(N)$ 개임을 관찰할 수 있다. 이러한 유용한 전이만을 갱신하는 것은 Heap 등의 자료구조를 통해서 가능하다.
A. 육각형 영역
서브태스크 2 (9점)
영역의 형태가 항상 정삼각형이다. 거리가 $d$ 인 점의 개수가 $d+1$ 개라는 점을 사용하면 문제의 답을 적절한 수식으로 표현할 수 있다. $\sum n, \sum n^2$ 를 공식을 사용하여 $O(1)$ 에 계산하면 전체 문제가 $O(1)$ 에 풀린다.
서브태스크 3 (20점)
육각 격자를 컴퓨터에서 다루는 가장 편한 방법은, 그냥 1/4 방향을 $y$ 축, 3/6 방향을 $x$ 축으로 둔 후, 2/5 방향을 $x=y$ 축으로 두는 것이다. 한쪽 축을 $x$ 축이 되도록 찌그러뜨린다고 생각하면 이해가 쉬울 것이다. 이제 격자 형태로 다룰 수 있으니 문제를 쉽게 해결할 수 있다. $(x, y)$ 에 인접한 셀은 $(x, y+1), (x+1, y+1), (x+1, y), (x-1, y), (x-1, y-1), (x, y-1)$ 이 6가지가 있다. $2000 \times 2000$ 격자에 주어진 도형의 외곽선을 적절히 그린 후, 내부를 Flood-fill로 채우고, BFS로 정답을 계산해 주자.
서브태스크 5 (47점)
$B = 0$ 이라 셀의 수만 구해주면 된다. 셀은 $(x, y)$ 를 좌하단 점으로 한 $1 \times 1$ 크기의 직사각형인데, 여기서는 그렇게 생각하지 말고 그냥 $(x, y)$ 에 있는 크기 0의 점이라고 생각하자. 이렇게 볼 경우, 셀의 자취 역시 다각형으로 볼 수 있다. 우리는 특정한 다각형 내부와 경계에 있는 점의 개수를 세어 주면 된다.
픽의 정리에 의하여, 다각형의 넓이가 $A$, 내부에 있는 점의 개수를 $i$, 다각형의 변 위에 있는 점의 수를 $b$ 라고 하면 $A = i + \frac{b}{2} - 1$ 이라는 식이 성립한다고 한다. $A$ 는 ccw로 쉽게 계산할 수 있고, $b = \sum L_i$ 이다. 고로 $i$ 를 구할 수 있다. 셀의 개수는 $i + b$ 이다.
서브태스크 6 (66점)
일반성을 잃지 않고 $d$ 의 합을 구하는 문제만 해결하자. $A$ 의 합은 서브태스크 5와 같이 픽의 정리를 사용하면 된다.
도형 안에 있는 격자점의 개수가 $L^2$ 에 비례하기 때문에, 단순한 방법으로 해결할 수 없다. 이를 해결하기 위해서 약간의 아이디어가 필요하다. 사실 여기가 이 문제에서 머리를 써야 하는 유일한 순간이다.
Claim (79brue). 다각형 상의 경로 $d$ 에서, $x$ 축 방향으로 이동한 횟수를 $a$, $x = y$ 축 방향으로 이동한 횟수를 $b$, $y$ 축 방향으로 이동한 횟수를 $c$ 라고 하자 ($+x, -x$ 축으로 움직였는지 여부는 신경쓰지 않고 모두 1회로 간주한다). 임의의 두 점을 잇는 모든 최단 경로에 대해서, $a + b + c$ 가 일정한 것은 자명하다. 이에 더해, 임의의 두 점을 잇는 모든 최단 경로에 대해서 $a, b, c$ 모두가 항상 일정하다.
Intuition? 엄밀한 증명은 모르지만 대략의 직관은 다음과 같다. 경로가 다각형의 밖으로 넘어가도 된다고 생각하면, 두 점을 잇는 최단 경로는 항상 $a, b, c$ 가 일정함을 쉽게 보일 수 있다. 실제로는 경로가 다각형의 안으로 들어가야 하는데, 처음으로 다각형을 벗어나는 지점에서 경로를 틀어야 하는 방향은 항상 고정되어 있다. 이를 계속 반복하면 결국 항상 $a, b, c$ 중 다각형이 강제하는 방향으로만 경로를 밀 수밖에 없고, 고로 항상 $(a, b, c)$ 쌍이 유일하게 최단 경로가 형성된다. 비슷한 내용이 79brue (반딧불 학생) 의 블로그에도 적혀 있다. 그 외에, 최단 경로가 저 쌍이 다르면 아마 둘을 조합해서 더 짧은 경로를 만드는 증명도 존재할 것 같다고 생각만 해 봤다.
여담으로 이 주장은 IOI 2012 이상적인 도시 문제의 그것과 동일한 주장이고, 그래서 접근도 대충 비슷하다. 처음으로 시도해 볼 수 있는 것은 모든 최단 경로에 대한 $a$ 의 합, $b$ 의 합, $c$ 의 합 등을 각각 구해서 더하는 방법이지만 이렇게는 잘 되지 않는다.
한편, 여기서 $a + b$ 의 합을 구하는 것을 생각해 보자. 이 경우, 가로로 움직일 경우 무조건 1의 비용이 들고, 세로로 움직일 경우 무조건 0의 비용이 든다. 0의 비용으로 움직일 수 있는 셀을 하나의 컴포넌트로 묶어주면 그 셀들이 $y$ 축에 평행한 줄을 이룰 것이다. 이 줄 각각을 정점으로 두고, 가로로 서로 움직일 수 있는 줄 간에 간선을 이어주면, 트리를 얻을 수 있다. 이 트리에서 $(0, 0)$ 을 포함하는 정점을 루트로 모든 정점에 대해서 (해당 정점의 깊이) * (해당 정점에 속하는 셀의 수) 를 다 합해주면 되는데, 이는 트리의 크기에 선형인 시간에 할 수 있다. 각 셀의 양 끝 경계는 입력으로 들어온 다각형의 경계선 중 하나이니, 셀의 개수는 많아야 $\sum L \le 200,000$ 개이다. 고로, 트리만 만들 수 있으면 $a + b$ 의 합을 구해줄 수 있고, 전체 도형을 계속 60도씩 Rotate해주면 $b + c$, $c + a$도 동일하게 알 수 있기 때문에 전체 문제가 해결된다.
최종적으로, 트리를 만드는 법을 알아보자. 각 $x = c$ 에 대해서, $(c, d)$ 와 $(c, d + 1)$ 이 다각형의 경계로 분리되는 모든 $d$ 를 모아주자. 이러한 $d$ 의 개수는 무조건 짝수개일 것이고, 이들을 두 개씩 묶어주면 모든 셀을 알 수 있다. 입력으로 주어진 다각형의 경계선을 따라가면서 이러한 $d$ 를 모든 $c$ 에 대해서 저장해 주면, 트리를 만들 수 있다는 것이다. 실제로는 경계선 처리 때문에 모든 $c$ 에 대해서 이것을 저장하는 게 그렇게 간단하지만은 않고, 이 문제 랑 동일하게 귀찮은 케이스 처리를 해 줘야 한다. 어려울 것은 없으나, 한번 실수를 했다면 인내심이 조금 많이 필요할 수도 있다. 시간 복잡도는 $O(\sum L \log \sum L)$ 이다.
만점 풀이
여기까지 왔으면 만점 풀이가 어떤 식일지는 뻔하다. 만약 셀 근처에 다각형의 정점이 없다면, 그 근처에서 셀의 경계를 이루는 선분의 집합은 동일할 것이다. 셀의 경계를 이루는 두 선분이 고정되면, 그 사이에 있는 정점들은 Path를 이루며, 각 정점에 포함된 셀의 개수가 등차수열을 이룬다. 이러한 연속된 정점들을 적당한 Super node의 형태로 묶어준 후, 이러한 연속성이 바뀌는, 모서리 근처 에서만 새로운 정점과 간선을 만들어 주면, 결국 하나의 정점 앞뒤로 $O(1)$ 개의 super node만이 생기게 되어, 트리의 크기를 $O(n)$ 으로 줄여줄 수 있고, 이 트리에서 우리가 원하는 값은 적당한 등차수열의 합 공식을 사용하여 잘 구할 수 있다.
고로, 선분의 집합을 std::set 에 넣고 스위핑하면서, 모서리 근처 에서 선분이 바뀌게 되면 super node를 적당히 넣고 빼어준 후 연결 관계를 조정해 주면 된다. 근처 라는 말의 미묘함을 다루는 과정이 아주 흥미로울 것으로 짐작된다. 66점 풀이의 예외 처리를 완벽하게 가져 오는 것이 중요하다.
시간 복잡도는 $O(n \log n)$ 이다.
B. 밀림 점프
서브태스크 1 (4점)
무조건 오른쪽으로 한 칸 갈 수 있고 그것만 할 수 있다. 답은 $C - B$ 이다.
서브태스크 2 (12점)
$x$ 에서 이동할 수 있는 두 나무를 $O(N)$ 시간에 찾을 수 있고, 전체 그래프를 $O(N^2)$ 시간에 찾을 수 있다.
이 그래프에서 Floyd-Warshall을 사용하여 모든 쌍 최단 경로를 계산하면, 각 쿼리에 대해 시작점 / 도착점을 모두 시도해 보는 것으로 $O(N^2)$ 시간에 답을 찾을 수 있다.
전체 시간 복잡도는 $O(N^3 + QN^2)$ 이다.
서브태스크 3 (25점)
각각의 쿼리에 대해서 $O(N)$ 시간에 문제를 해결한다. BFS를 하는데, 처음 시작할 때 $A, A+1, \ldots, B$ 에 있는 모든 정점들을 큐에 넣고 거리를 0으로 한다. 이렇게 탐색을 하면, 모든 점에 대해서 ($[A, B]$ 의 어떤 점에서 출발했을 때 최단 경로) 를 계산할 수 있다. 자세한 설명은 종만북을 참고하면 좋다.
그래프의 간선 수가 $O(N)$ 이니 전체 시간 복잡도는 $O(N^2 + QN)$ 이다.
서브태스크 4 (37점)
스택을 사용하면 그래프를 $O(N)$ 시간에 찾을 수 있다.
일반성을 잃지 않고, $H[y] > H[x]$ 이자 $x$ 보다 작은 최대 $y$ 만 찾아주자. $H[x - 1], H[x - 2] \ldots$ 를 순서대로 보면, 이 중 몇가지 원소들은 지금까지 본 원소 중 가장 큰 (최댓값을 갱신하는) 원소일 것이다. $x$ 에 대한 답을 찾기 전에, 이러한 원소들을 스택에 저장하자. 즉, $H[i]$ 가 스택에 있다면, $H[i+1], H[i+2] \ldots H[x-1]$ 이 모두 $H[i]$ 보다 작다는 뜻이다. 스택에는 이러한 $H[i]$ 가 $i$ 가 증가하는 순서대로 bottom에서 top으로 저장되어 있을 것이다.
만약 이러한 원소들을 스택에 저장했다면, 스택에 있는 원소들만이 $y$ 가 될 수 있는 유일한 후보이다. 스택에 없다면, 이보다 큰 원소가 $x$ 에 더 가까운 위치에 있다는 뜻이기 때문이다. 정확히는, 스택에서 $H[x]$ 보다 큰 값을 가지는 첫 $y$ 가 답이 될 것이다.
$H[0], H[1], \ldots, H[x-1]$ 에 대해서 이 스택을 관리했다면, $H[0], H[1], \ldots, H[x]$ 에 대해서도 이 스택을 관리할 수 있다. 단순히 현재 스택의 top에서부터 $H[x]$ 보다 작거나 같은 원소들을 제거해 주면 된다. 제거가 끝나는 시점에 스택의 top에 있는 원소는 $H[y] > H[x]$ 를 만족하는 첫 $y$ 일 것이다. 그래서 그 $y$ 에 간선을 이어주면 된다. 간선을 이어준 후, 그냥 $H[x]$ 를 스택에 넣어주면 된다.
서브태스크 5 (60점)
서브태스크 5부터는 그래프를 단순히 탐색하는 것으로는 불가능하다. 그래프의 성질을 사용하여 답을 찾는 쉬운 알고리즘을 고안한 후, 자료구조를 사용하여 최적화하자.
$L[v], R[v]$ 를 $v$에서 왼쪽 / 오른쪽으로 움직였을 때 다음 위치라고 하자. 만약 그러한 $L[v], R[v]$ 가 없으면 해당 값은 $v$ 로 둔다. 아래와 같은 관찰을 할 수 있다.
Theorem. $H[L[s]] \le H[t], H[R[s]] \le H[t]$ 이면, $H[L[s]], H[R[s]]$ 중 높이가 높은 방향으로 가는 것이 항상 최적이다.
Proof. 일반성을 잃지 않고 $H[L[s]] < H[R[s]]$ 라고 한다. 또한, $H[L[s]] = H[t], H[R[s]] = H[t]$ 인 경우 자명히 참이니, $H[L[s]] < H[t], H[R[s]] < H[t]$ 임을 가정하자. 귀류법을 사용한다. 최적해가 $L[s]$ 로 가야만 하는 경우를 가정하자. $L[s]$ 로 갔다가 $R[s]$ 로 돌아오면 그냥 $R[s]$ 로 가는 것이 나으니, $L[s]$ 쪽으로 어느 정도 가다가 $H[R[s]]$ 의 높이를 처음으로 뛰어넘는 위치 $x$ 에 언젠가 도달할 것이다. 한편, $R[s]$ 에서는 항상 이 $x$ 로 바로 움직일 수 있다. 고로 이 경로를 $R[s] \rightarrow x$ 로 대체하면 $L[s]$ 를 거치지 않는 다른 최적해를 찾을 수 있다. $\blacksquare$
고로, 위 조건이 만족할 때까지 높이가 높은 방향으로 움직이면 된다. 이 조건이 만족하지 않는다는 것은, 더 이상 움직일 곳이 없거나, 높은 곳으로 움직였을 때 $H[t]$ 를 초과한다는 뜻이다. 전자의 경우에는 종료하고, 후자의 경우에는 높이가 낮은 곳으로 움직이면 된다. (낮은 곳으로 움직일 경우, 높은 곳의 위치가 계속 동일하게 유지되어서, 다음에 높은 곳으로 다시 움직여야 할 상황이 발생하지 않는다.)
이렇게 하여 $O(N)$ 시간에 답을 찾는 그리디 알고리즘을 찾을 수 있다. 이는 Sparse table로 최적화할 수 있다. $v$ 에서 $2^k$ 번 높은 곳으로만 / 낮은 곳으로만 움직였을 때의 위치 $High[k][v], Low[k][v]$를 $O(N \log N)$ 시간에 계산해 두면, 각 쿼리마다 이분 탐색으로 $O(\log N)$ 시간에 답을 계산할 수 있다.
서브태스크 6 (81점)
$B, L[B], L[L[B]] \ldots$ 를 타고 따라가면 값이 증가하는 체인을 찾을 수 있다. 이 체인 상에 있는 모든 점 $v$ 는 $H[v] > max(H[v+1], \ldots, H[B])$ 를 만족하고, 이 체인 상에 없는 모든 점 $v$ 는 이를 만족하지 않는다. $[C, D]$ 구간에 도달하기 위해 오른쪽으로 가려면 무조건 해당 체인에 있는 정점을 거쳐야 한다는 것이다. 즉, 체인 상에 있는 정점만 시작점으로 고려하면 된다.
체인에 있는 점 중 높이가 $H[C]$ 이하인 가장 높은 점을 고르자. 이 점에서 $C$ 에 도달할 수 없다면, 이 점과 $C$ 사이에 높이가 $H[C]$ 초과인 정점이 있다는 것이니 어느 점에서 시작해도 도달이 불가능하다. 한편, 이 점이 아닌 점에서 $C$ 에 도달하는 경로는, 모두 이 점에서 도달하는 것보다 길이가 짧거나 같다. 그래서, 체인에 있는 점 중 $H[v] < H[C]$ 인 가장 높은 점 $v$ 에서 시작하면 된다. 이러한 점은 Sparse table로 찾을 수 있고, 이후 서브태스크 5처럼 진행하면 된다.
만점 풀이
$[B, C-1]$ 구간에서 높이가 최대인 점을 $h$ 라고 정의하자. 모든 가능한 이동은 이 $h$ 를 거치거나, 아니면 $h$ 를 거치지 않고 그 위로 올라간다.
$h$ 를 거치는 경우, $h$ 에서 한 번의 오른쪽 점프로 바로 $[C, D]$ 로 도달할 수 있어야 한다. 만약 그렇지 않은 경우, 목적지에 도달할 수 있는 방법은 없다. 한 번의 점프로 $[C, D]$ 에 도달할 수 있는지는 쉽게 판단할 수 있다. $[A, B]$ 에서 $h$ 에 도달하는데 필요한 점프는 서브태스크 6의 방법으로 구할 수 있다. 이 때 $h$ 에 도달하기 위해 시작한 지점이 있을 텐데, 이 지점을 $g$ 라고 하자. 즉, $g$ 는 $B$ 에서 시작하는 체인에서 $H[g] < H[h]$ 를 만족하는 첫 $g$ 이다.
$h$ 의 위로 올라갈 경우를 고려해 보자. 일단 첫번째로, $g$ 혹은 $g$ 보다 오른쪽의 위치에서 출발하는 경우는 이미 위 케이스에서 고려가 되기 때문에 ($g$ 에서 출발하는 것이 더 이득이기 때문) 여기서 고려해 줄 필요가 없다. 즉 시작점이 $L[g], L[L[g]] \ldots $ 라고 가정할 수 있다는 것인데, 이 위치들은 모두 $h$ 보다 높이가 높다. 고로 여기서 오른쪽으로 점프했을 경우, $h$ 를 거치지 않고, $C$ 혹은 그 이후 위치에 떨어지게 된다. 이렇게 될 경우 높이가 높아질 수록 $[C, D]$ 구간에서 벗어나는 게 손해이기 때문에, $L[g]$ 위치에서 오른쪽으로 한 번 뛰어서 $[C, D]$ 에 도착하는지만 확인해 주면 된다.
$[B, C-1]$ 구간의 최대는 Segment tree로 계산할 수 있지만, 사실 있던 sparse table을 재활용해서 계산하는 것이 편하다. 어떻게 할 수 있는 지는 연습으로 남긴다. 최종 시간 복잡도는 $O((N + Q) \log N)$ 이다.
C. 도로 폐쇄
서브태스크 1 (5점)
모든 간선에 0번 정점에 인접해 있어서, 각 정점의 차수가 최대 $k$ 여야 한다는 것은 그냥 그래프의 간선이 최대 $k$ 개여야 한다는 것과 동일하다. 단순 정렬로 해결할 수 있다.
서브태스크 2 (12점)
그래프가 Path를 이룬다. 이미 모든 정점의 차수가 2 이하라, $k \ge 2$ 인 경우에는 그냥 그래프 전체를 가져갈 수 있다. $k = 1$ 일 경우에는 고른 간선이 인접해서는 안된다. $0 \ldots i$ 번 정점까지만 고려했을 때, 인접하지 않게 고를 수 있는 최대 간선의 가중치 합을 $dp[i]$ 라고 두면, 단순한 점화식을 유도할 수 있다.
서브태스크 4 (36점)
$k$ 를 고정할 경우 문제를 트리 DP로 해결할 수 있다.
$DP[v][0]$ 을, $v$ 를 루트로 한 서브트리에서 각 정점의 차수를 $k$ 이하로 유지하기 위해 지워야 할 간선의 최소 가중치 합이라고 두자. $DP[v][1]$ 도 동일하게 두는데, $v$ 와 $v$ 의 부모를 잇는 간선을 유지한다고 가정하고 두자. 즉, $DP[v][1]$ 의 경우 $v$ 의 차수는 $k-1$ 이하여야 한다 (부모를 잇는 간선을 추가할 것이라서).
$v$ 의 자식 $w$를 고를 때, 현재 정점과 자식을 잇는 간선을 고른다면 $DP[w][1]$ 의 비용이, 현재 정점과 자식을 잇는 간선을 고르지 않는다면 $DP[w][0] + cost(v, w)$ 의 비용이 필요하다. 첫 번째 옵션은 최대 $k$ 개 혹은 $k-1$ 개 고를 수 있다. 기본적으로 $DP[w][0] + cost(v, w)$ 를 고른다고 하고, $min(DP[w][1] - DP[w][0] - cost(v, w), 0)$ 중 가장 작은 $k$ 개 혹은 $k-1$ 개를 골라서 이 값을 최대한 줄여주는 전략을 사용할 수 있다.
이를 사용하면 각 정점마다 $deg(v)$ 개의 원소를 정렬해야 하니 하나의 $k$ 에 대해서 문제가 $O(n \log n)$ 에 해결된다. 이를 $n$ 번 반복하면 $O(n^2 \log n)$ 에 문제를 해결할 수 있다.
서브태스크 5 (53점)
정해에는 크게 도움이 되지 않는 풀이이다.
고정된 $k$ 에 대해서, 차수가 $k$ 초과인 정점들만 모아서 induced subgraph를 만들자. 이 induced subgraph의 각 forest에 대해서 문제를 해결할 것이다.
각 Forest에 대해서 그리디하게 문제를 해결하자. Forest를 구성하는 간선들은, 양 끝 정점의 차수가 모두 $k$ 를 초과하니, 이들을 최대한 많이 사용해서 차수를 줄여주는 것이 좋을 것이다. 리프에서부터 Bottom-up으로, 만약 현재 정점의 차수가 $k$ 를 초과한다면 부모로 가는 간선을 끊어주는 것을 반복한다. 이 과정을 완료해도 차수가 $k$ 가 초과하는 정점들은 있을 것인데, 이러한 정점의 모든 간선은 현재 차수가 $k$ 를 초과하지 않는 정점으로만 이어져 있다. 고로 무조건 하나씩 지워서 차수를 줄여야 하니 $max(deg(v) - k, 0)$ 을 단순히 더해주면 된다.
시간 복잡도가 중요한데, 정점 $v$ 가 고려되는 $k$는 $0$ 이상 $deg(v) - 1$ 이하 뿐이다. 고로 모든 $k$ 에 대해서 고려하는 정점의 개수를 합하면 $O(n)$ 이 된다. 고로 전체 시간 복잡도는 $O(n)$ 이고, 필요하면 $O(n \log n)$ 정도로 구현해도 괜찮을 것이다.
만점 풀이
만점 풀이는 약간 테크니컬하다. High-level idea만 말하자면, DP 식을 적당히 변형하고 적절한 자료구조를 사용하여, $deg(v) > k$ 인 점들에 대해서만 DP 값을 업데이트해도 되게끔 방법을 바꾸는 것이다. 각 점에 대해서 DP 값을 업데이트 하는데 $O(\log n)$ 시간이 든다면, 위에서와 같이 정점 $v$ 가 고려되는 $k$는 $0$ 이상 $deg(v) - 1$ 이하 뿐이니 $O(n \log n)$ 풀이를 얻을 수 있다.
$diff(v) = DP[v][1] - DP[v][0] - cost(v, par(v))$ 라고 정의하자. 이렇게 정의할 경우 점화식은 다음과 같다.
- $DP[v][0] = \sum_{w \in child(v)} (DP[w][0] + cost(v, w)) + $ ($diff(w)$ 가 가장 작은 최대 $k$ 개의 원소 합)
- $diff(v) = -cost(v, par(v))$ - ($k$ 번째로 큰 $diff(w)$)
그리고 답은 $DP[0][0]$ 이다. 이 점화식은 서브태스크 4의 점화식의 식을 적당히 비튼 것에 불과하다.
여기서 두 가지 관찰을 할 수 있는데, 첫 번째로 $DP[v][0]$ 에 대한 생각을 아예 할 필요가 없다는 것이다. 결국 $DP[0][0]$ 은 모든 정점 $v$ 에 대해서 $\sum_{w \in child(v)} cost(v, w) + $ ($diff(w)$ 가 가장 작은 최대 $k$ 개의 원소 합) 을 합한 것에 불과하다. 첫 번째 항은 모든 간선의 가중치 합이고, 두 번째 항은 완전히 $diff(w)$ 로만 표현할 수 있다. 고로 이제부터는 $diff(v)$ 만 관리하면 된다.
두 번째로, $deg(v) < k$ 일 때 $diff(v)$, 그리고 $diff(w)$ 가 가장 작은 $k$ 개의 원소 합은 변하지 않는다는 점이다. 정확히 말해서, $k$ 개의 원소 합 자체는 변할 수 있으나, 그건 $diff(w)$ 가 변할 때 알아서 부모의 자료 구조를 잘 바꿔주면 된다.
여기까지 왔으면 전체 문제를 해결할 준비는 모두 끝났다. $k = N - 1, N - 2, \ldots, 1$ 로 감소시키면서 문제를 해결하자 ($k = 0$ 은 예외처리했다). 우리가 관리해야 할 것은 각 정점에 대해서 $diff(w)$ 의 가장 작은 $k$ 개 원소들의 합과, $diff(w)$ 그 자체이다. $k$ 가 하나씩 감소한다면
- $deg(v) > k$ 인 원소들에 대해서 $diff(w)$ 의 $k+1$ 번째 원소를 빼줘야 하며
- $deg(v) \geq k - 1$ 인 원소들에 대해서 $diff(v)$ 를 갱신하고 부모 자료구조에 업데이트 해야 한다. 이후 부모에서도 이를 반복한다.
두 번째 과정 때문에 루트로 가는 모든 정점을 다 갱신해야 하는 것이 아닐까 생각할 수 있지만, $deg(v) \le k - 2$ 인 경우에는 $diff(v)$ 가 상수라서 갱신할 필요가 없다. 고로 DFS preorder의 역순으로 돌리면서 위 과정을 반복만 해 주면 그만이다.
위 연산들을 지원하기 위해서는 다음과 같은 자료구조가 필요하다:
- 삽입
- 삭제
- $k$ 번째 원소 찾기
- 가장 작은 $k$ 개 원소 합 찾기
- $k$ 를 1씩 줄이기
이를 구현하는 가장 간단한 방법은 std::set 을 두 개 두는 것이다. 첫 번째 set은 $k$ 개의 가장 작은 원소들을 가지고 있고, 두 번째 set은 그 외 다른 원소들을 가지고 있다. 이 때 첫 번째 set의 크기가 $k$ 를 넘어가면 두 번째로 넘겨주고, $k$ 를 밑돌면 두 번째에서 빼주는 식으로 위 연산들을 모두 worst case $O(\log N)$ 시간에 구현해 줄 수 있다.
이제 각 노드에 대해서 위와 같은 자료구조를 선언하고, $k$ 가 감소할 때 마다 이 자료구조를 위에서 설명한 대로 잘 갱신해 주면 된다.