この問題はいくつかの解法がありますが、この解説では俗に Auxiliary Tree と呼ばれるテクニックを紹介します。

解法の方針

はじめにこの問題を解く方針を説明します。
次数 1 の頂点に塗られた色 \(c\) を固定します。すると、固定したときの問題は木 DP で \(\mathrm{O}(N)\) で解くことが出来ます。(これは容易ではないですが青 diff 程度だと思うので分からない方は練習問題として考えてみましょう) しかし、これを全ての色に対して計算すると \(\mathrm{O}(N^2)\) かかり TLE してしまいます。
色 \(c\) を固定したときの問題が \(\mathrm{O}(N)\) かかる理由は、木の頂点が \(N\) 個あることによります。しかし、この \(N\) 個の頂点の中には、例えば絶対に使用されない頂点のような本質的に要らない頂点がたくさんあり、色 \(c\) に関して DP するときにこのような頂点も含めて計算することは計算量の無駄です。そこで木をうまく圧縮することを考えます。
木を圧縮する際に残す頂点を考えます。色 \(c\) の頂点はもちろん残すとして、頂点同士の位置関係を失わないためには、色 \(c\) である頂点同士の LCA となりうる頂点も残すと良さそうです。(ここで木 \(T\) は頂点 \(1\) を根とする根付き木とみなして考えています。) このようにして必要な頂点を選んで、子孫関係のある部分に辺を貼ると左図から右図のようになります。右の図で得られた木を元に DP を計算しても問題なさそうです。

(図：頂点 \(6, 7, 8, 14, 15\) が色 \(c\) で塗られている時の圧縮の仕方)

このようにして木を圧縮して出来る木を 指定された頂点たちの最小共通祖先関係を保って木を圧縮してできる補助的な木 と呼びます。(俗に Auxiliary Tree, virtual tree, 圧縮木 と呼ばれることがあります。指定された… は長いので以降では Auxiliary Tree と呼びます。)

• 命名に関してはこちらの記事も参考にしてください。　Auxiliary Tree に関する Kyopro Encyclopedia of Algorithms の記事

Auxiliary Tree の性質

Auxiliary Tree をより形式的に説明すると次のようになります。 根付き木 \(T\) および頂点集合 \(X\) が与えられたとき、頂点集合 \(A(X)\) を

\[A(X) = \lbrace \mathrm{lca}(x, y) \vert x \in X, y \in X \rbrace\]

と定義して、\(A(X)\) に含まれる頂点の間に \(T\) での子孫関係のある部分に辺を張ってできる根付き木を Auxiliary Tree と呼びます。

Auxiliary Tree の性質として次の性質が挙げられます。

性質 1

\(\vert A(X) \vert \leq 2 \vert X \vert - 1\) である。

(証明) 木 \(T\) に対して DFS の行きがけ順を 1 つ自由に取る。そして、\(X\) の頂点を行きがけ順に登場する順番に並べたものを \(x_1, x_2, \dots, 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)\) は

\[ \begin{aligned} &A(\lbrace x_1, \dots, x_m\rbrace) \\ &= A(\lbrace x_2, \dots, x_m\rbrace) \cup \lbrace x_1, \mathrm{lca}(x_1, x_2) \rbrace \end{aligned} \]

と表現できることが確認できる。この式から

\[\vert A(X)\vert \leq \vert A(\lbrace x_2, \dots, x_m\rbrace)\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\) 頂点になることが証明できればよく、\(T\) が完全二分木で \(X\) が葉の頂点からなる頂点集合である場合を考えると \(\vert A(X) \vert = 2 \vert X \vert - 1\) となるケースを構成できます。

この性質は、Auxiliary Tree の頂点数は \(\mathrm{O}(|X|)\) 、つまり元となる頂点集合のサイズの定数倍で抑えられることを意味します。元の問題の解法ではこの事実を利用して計算量削減を行っています。

Auxiliary Tree の構築

頂点集合 \(X\) に関する Auxiliary Tree の構築は1 回の生成あたり \(\mathrm{O}(|X| (\log |X| + \log N))\) で行うことが出来ます。(\(N\) は元の根付き木の頂点数。また、木 \(T\) に関する前計算が \(\mathrm{O}(N \log N)\) 程度必要です)

構築方法を説明します。次の性質を利用します。

性質 2

木 \(T\) に対して DFS の行きがけ順を 1 つ自由に取る。そして、\(X\) の頂点を行きがけ順に登場する順番に並べたものを \(x_1, x_2, \dots, x_m\) とする。
このとき \(A(X)\) に関して次の事実が成り立つ。

\[A(X) = X \cup \lbrace \mathrm{lca}(x_i, x_{i+1}) \vert 1 \leq i \lt m \rbrace\]

(証明)

性質 1 の証明で示した

\[ \begin{aligned} &A(\lbrace x_1, \dots, x_m\rbrace) \\ &= A(\lbrace x_2, \dots, x_m\rbrace) \cup \lbrace x_1, \mathrm{lca}(x_1, x_2) \rbrace \end{aligned} \]

という式を利用すると

\[ \begin{aligned} &A(\lbrace x_1, \dots, x_m\rbrace) \\ &= A(\lbrace x_2, \dots, x_m\rbrace) \cup \lbrace x_1, \mathrm{lca}(x_1, x_2) \rbrace \\ &= A(\lbrace x_3, \dots, x_m\rbrace) \cup \lbrace x_1, \mathrm{lca}(x_1, x_2), x_2, \mathrm{lca}(x_2, x_3) \rbrace \\ &\vdots \\ &= A(\lbrace x_m \rbrace) \cup \lbrace x_1, \mathrm{lca}(x_1, x_2), x_2, \dots, x_{m-1}, \mathrm{lca}(x_{m-1}, x_m) \rbrace \\ &= \lbrace x_1, \mathrm{lca}(x_1, x_2), x_2, \dots, x_{m-1}, \mathrm{lca}(x_{m-1}, x_m), x_m \rbrace \\ \end{aligned} \]

となることがわかる。よって示された。(証明終わり)

性質 2 を利用すると、以下の手順で \(A(X)\) に登場する頂点を行きがけ順に並べたものを得ることが出来ます。

• 前計算の段階でダブリングや HLD を用いて木 \(T\) の任意の 2 頂点間の LCA を高速に計算できるようにする。
• \(X\) に登場する頂点を行きがけ順に並べて、順に \(x_1, x_2, \dots, x_m\) と呼ぶ。
• \(\mathrm{lca}(x_1, x_2), \dots, \mathrm{lca}(x_{m-1}, x_m)\) を計算する。
• \(x_1, x_2, \dots, x_m, \mathrm{lca}(x_1, x_2), \dots, \mathrm{lca}(x_{m-1}, x_m)\) を行きがけ順にソートして重複を取り除く。

登場する頂点が全て分かれば、あとはこの列を利用して stack を用いながら頂点を走査していくことで Auxiliary Tree を構築できます。(詳細は略します。詳しい手順は yaketake08 さんのブログ記事 の 「ソート2回の方法」の項を参考にしてください。)

さて、以上の内容を踏まえて、各色ごとにその色が塗られた頂点の集合から Auxiliary Tree を生成して木 DP することで、この問題を \(\mathrm{O}(N \log N)\) で解くことが出来ます。これは十分高速です。