Colorful Ancestors

Consider the contribution of each color independently.

Fix a color $c$. Every vertex $v$ has a value $T_c(v)$, which is either $0$ or a vertex of color $c$. The contribution of color $c$ to the answer is the number of ordered pairs with equal $T_c$ values, namely

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

where $cnt_c(x)$ is the number of vertices $v$ such that $T_c(v)=x$. Here, $x$ is either $0$ or a vertex of color $c$.

Now consider a vertex $x$ of color $c$. A vertex $v$ satisfies $T_c(v)=x$ if and only if

• $v$ lies in the subtree of $x$.
• On the path from $x$ down to $v$, excluding $x$ itself, there is no vertex of color $c$.

So we start with all vertices in the subtree of $x$, then subtract the subtrees of the closest same-color descendants of $x$.

Thus,

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

where the sum is over vertices $y$ such that $x$ is the closest same-color ancestor of $y$.

It remains to compute, for every vertex $y$, its closest same-color ancestor $sameParent(y)$. This can be done in one DFS from the root. During DFS, maintain $last[c]$, the deepest vertex of color $c$ on the current root-to-current path.

When entering a vertex $v$, the current value of $last[C_v]$ is exactly $sameParent(v)$. Store it, then set $last[C_v]=v$. When leaving $v$, roll it back. Each vertex is processed once, so this is linear.

We also need the number of vertices with $T_c(v)=0$. For a fixed color $c$, consider vertices of color $c$ that have no same-color ancestor. Their subtrees are disjoint, and every vertex inside them has some color-$c$ ancestor. Therefore,

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

where the sum is over color-$c$ vertices $x$ with no same-color ancestor.

In the implementation:

• If $sameParent(v)=0$, add $subtree(v)$ to $rootSum[C_v]$.
• Otherwise, add $subtree(v)$ to $sameChildSum[sameParent(v)]$.

Then the answer is

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

Since $N$ can be $2\,000\,000$, a recursive DFS may cause stack overflow. An iterative DFS is safer.

The time complexity is $O(N)$, and the memory complexity is $O(N+K)$.