APIO 2021

원래 한 150점 어치 풀이를 적어놓은 게 있었는데, 포맷하다가 글이 사라져서 더 이상 글을 쓰고 싶지 않았다. 그래도 아예 안 적을 수는 없으니 늦게나마 글을 작성해 본다.

문제 풀이 한 줄 요약

• 육각형 영역: 다각형 내부의 모든 격자점 간의 거리 합을 구하는 문제로 환원할 수 있다. 다각형의 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$ 가 감소할 때 마다 이 자료구조를 위에서 설명한 대로 잘 갱신해 주면 된다.