알록달록한 트리

정답을 색별로 나누어 생각하자.

고정된 색 $c$에 대해, 모든 정점 $v$는 $T_c(v)$라는 값을 가진다. 이 값은 $0$이거나 색이 $c$인 어떤 정점이다. 색 $c$가 전체 답에 기여하는 값은 같은 $T_c$ 값을 가지는 순서쌍의 개수이므로,

$$\sum_x cnt_c(x)^2$$

이다. 여기서 $cnt_c(x)$는 $T_c(v)=x$인 정점 $v$의 개수이다. $x$는 $0$ 또는 색이 $c$인 정점이다.

이제 각 색에 대해 이 값을 빠르게 계산하면 된다.

색이 $c$인 정점 $x$를 보자. $T_c(v)=x$인 정점 $v$는 다음 조건을 만족한다.

• $v$는 $x$의 서브트리 안에 있다.
• $x$에서 $v$로 내려가는 경로에서 $x$를 제외하고 색이 $c$인 정점이 없다.

즉, 처음에는 $x$의 서브트리 크기만큼 후보가 있고, 그중 $x$ 바로 아래의 같은 색 정점들이 담당하는 서브트리를 빼면 된다.

여기서 ``바로 아래의 같은 색 정점''이란, 색이 $c$인 정점 $y$ 중 $x$가 $y$의 가장 가까운 색 $c$ 조상인 정점을 뜻한다.

따라서

$$cnt_c(x)=subtree(x)-\sum subtree(y)$$

이다. 합은 $x$가 가장 가까운 같은 색 조상인 정점 $y$들에 대해 취한다.

남은 것은 각 정점 $y$에 대해 가장 가까운 같은 색 조상 $sameParent(y)$를 구하는 것이다. 이는 루트에서 DFS를 하면서 현재 경로 위에서 각 색의 가장 낮은 정점을 저장하면 된다.

DFS로 정점 $v$에 들어갈 때, 현재 $last[C_v]$가 $v$의 가장 가까운 같은 색 조상이다. 그 값을 $sameParent(v)$에 저장하고, $last[C_v]=v$로 바꾼다. DFS에서 $v$를 빠져나올 때는 원래 값으로 롤백한다. 각 정점마다 한 번씩만 처리하므로 선형 시간이다.

색 $c$에 대해 $T_c(v)=0$인 정점 수도 필요하다. 색 $c$인 정점 중, 같은 색 조상이 없는 정점들을 생각하자. 이 정점들의 서브트리는 서로 겹치지 않고, 그 안에 있는 정점들은 색 $c$ 조상을 하나 이상 가진다. 따라서

$$cnt_c(0)=N-\sum subtree(x)$$

이다. 합은 같은 색 조상이 없는 색 $c$ 정점 $x$들에 대해 취한다.

구현에서는 다음 값을 모은다.

• $sameParent(v)=0$이면 $rootSum[C_v]$에 $subtree(v)$를 더한다.
• $sameParent(v)\ne 0$이면 $sameChildSum[sameParent(v)]$에 $subtree(v)$를 더한다.

그 후 정답은

$$\sum_{c=1}^{K}(N-rootSum[c])^2+ \sum_{v=1}^{N}(subtree(v)-sameChildSum[v])^2$$

이다.

DFS는 재귀로 구현하면 $N=2\,000\,000$에서 스택 오버플로가 날 수 있으므로, 반복문 기반 DFS를 사용하는 것이 안전하다.

시간 복잡도는 $O(N)$, 메모리 복잡도는 $O(N+K)$이다.