https://kanpurin.hatenablog.com/entry/2020/04/21/154204

問題概要

${1\le N\le 2×10^5}$ 頂点の木が与えられ, 頂点 ${i}$ には色 ${1\le c_i\le N}$ が塗られている. ${k=1,2,\ldots,N}$ に対して, 以下の問題を解け.

色 ${k}$ が塗られている頂点を一度以上通るような単純パスの数を求めよ.

思考の流れ

一度以上通るっていうのは難しいので余事象「色 ${k}$ が塗られている頂点を一度も通らないような単純パスの数」を求めて全体から引く. 木の性質から, 頂点数 ${N}$ の木の単純パスの数は ${\frac{N(N+1)}{2}}$ となる. 木上の色 ${k}$ が塗られている頂点で木を分割すれば, グラフは森になり, 各連結成分の頂点数を求めればいい.

${\downarrow}$

色 ${k}$ で分割したときの単純パスの総数は適当な頂点からDFSすることで ${O(N)}$ で求めることができる. しかし, それだと ${N}$ 色すべてで ${O(N)}$ かかるので結局 ${O(N^2)}$ となり間に合わない.

${\downarrow}$

各頂点に対してその頂点の色以外を考えるときの更新に無駄が多そう.
DFSをする際にどういう風に更新していくか.
ここで, 図1のような状況を考える. 根の色を ${k}$ とし, 赤のところが同じ色 ${k}$ で塗られていると考える. 頂点の下の三角はその頂点の部分木である.

f:id:kanpurin:20200421144708p:plain — 図1

赤頂点の部分木(黒線三角)の連結成分を無視すると, 連結成分が2つあるが, 左側は図2のように求められる.

f:id:kanpurin:20200421145502p:plain — 図2

緑の三角で表された部分木のサイズ(頂点数) ${-}$ 水色の三角で表されたすべての部分木のサイズの和
とすると, 連結成分の頂点数を求めることができることがわかる.

緑色の三角形は根の子の部分木のサイズであるのですぐ求められるが, 水色の三角形はどのように求めるか.

${\downarrow}$

DFSで上に戻るとき, 色ごとに頂点を記録することが考えられるが, これだと上書きが面倒だったり頂点が多くなりすぎたりして難しい.
本来必要なものは水色の部分木のサイズの和のみであるのでそれを記録していくことを考える. こうすると図3の矢印のように根にDFSで根に帰ってきたときは色 ${k}$ の頂点の部分木(図2における水色の三角形)のサイズの和が求められているので, 色 ${k}$ による分割後の連結成分のサイズを求めることができる.

f:id:kanpurin:20200421151244p:plain — 図3

DFSで各色の頂点数を親に返しているわけではなく, グローバルにおいてる値(色 ${k}$ におけるこの値を ${X_k}$ とする)を更新していってるので, 各色について頂点数を加算して, それを色の頂点数としてしまうと, 図4のような状態を考えたとき, 赤色の頂点3つ分の部分木の頂点数を加算してしまっていて, 右側の連結成分が求められない.

f:id:kanpurin:20200421152323p:plain — 図4

ここで, 根の右側の子へDFSするときに ${X_k}$ の値を記録しておく(この値を ${Y}$ とする). こうすることで根に帰ってきたとき, 現在の ${X_k}$ と ${Y}$ の差を見ることで頂点数の和を求めることができるようになる.
また, 図5のようにDFSで親に戻るとき部分木の頂点数を ${X_k}$ に足す必要があるが, 部分木内の他の色 ${k}$ の部分木の頂点数をすでに足しているので余分に足してしまうことになる.