実行時間が短い提出が色ごとに抽出せずに処理しているように見えたので、これを参考に改めて考えました。その提出のプログラムの原理は詳しく確認していません。

Auxiliary Tree が不要なパターン

Auxiliary Tree （あるいは、元の木）は、ある部分木をとったときに、特定の色 \(c\) がついた頂点を \(2\) 頂点以上含み、それらがすべて葉となるような部分があります。また、全体の木はこの条件を満たす極大な部分に分解されます。このような部分木を (色 \(c\) の) block と呼ぶことにします。

色 \(c\) について考えるとき、もし block 内部の（つまり、色 \(c\) でない頂点からなる）部分グラフを \(1\) 頂点に圧縮しても問題なければ、具体的に Auxiliary Tree を構築しなくてもよい可能性が高いです。

今回の問題では、 block を \(1\) つ考えて、 block の葉の選び方を決めたとき、 block 内部の選び方は \(1\) 通りに決まるので、実は block 内部の構造は無視できます。

DP テーブルの持ち方

頂点 \(1\) を根として根付き木で考えます。

block の頂点集合 \(C\) のうち、最も根に近いものを選んで \(q\) とします。この \(q\) が葉である（色 \(c\) である）場合を考えます。 \(C\) から \(q\) を除いて、また最も根に近いものを \(r\) とします。このとき、 \(r\) を決めると block は高々 \(1\) 個に決まります。なぜなら、 block の色は \(r\) の親の色であり、 block の範囲は \(r\) から探索すれば決まるからです。

\(q\) が葉でない場合は \(q=1\) であって、各色ごとに高々 \(1\) つなので、

• \(q\neq 1\) の場合、 \(r\) ごとに高々 \(1\) 個に決まる block に \(r\) 番、
• \(q=1\) で色 \(c\) の block に \(N+c\) 番

を割り当てれば各 block について要素を持つテーブルが成立します。

\(c\) を \(1\) つ選ぶと、色 \(c\) の block すべての位置関係は根付き木で表せます。よって、各 block について親にあたる block を予め求めておくことで、さっきのテーブルを用いた DP ができます。

block の計算

次のテーブルを計算しながら DFS すればよいです。具体的な手順は提出を参考にしてください。

• dfs_stack : 根から現在探索している頂点までのパス。
• colored_anc[c] : 根から現在探索している頂点までのパスのうち、色 \(c\) が塗られている頂点のうち最も深いもの。存在しない場合は \(N+c\) 。
• same_colored_anc[v] : \(v\) と同じ色の頂点であって \(v\) の祖先であるもの（ただし、 \(v\) 以外）のうち、最も \(v\) に近いもの。存在しない場合は \(N+A _ v\) 。
• same_colored_anc_before[v] : \(v\) から same_colored_anc[v] に向かって進んだとき、 same_colored_anc[v] の直前に訪れるもの。 same_colored_anc[v]\({}\gt N\) の場合は same_colored_anc[v] と同じ。

計算量は \(O(N)\) です。

提出

ソースコードを読む場合は、まず main 関数のみ読むことを勧めます。 main 関数は 320 行目からです。ヘッダーを使っていて意図が分かりにくいと思われる部分はコメントを入れています。

提出 #50213569

このアプローチが有効な他の問題

• yukicoder