問題はこちら
問題概要
頂点の木が与えられ, 頂点には色が塗られている. に対して, 以下の問題を解け.
色が塗られている頂点を一度以上通るような単純パスの数を求めよ.
思考の流れ
一度以上通るっていうのは難しいので余事象「色が塗られている頂点を一度も通らないような単純パスの数」を求めて全体から引く. 木の性質から, 頂点数の木の単純パスの数はとなる. 木上の色が塗られている頂点で木を分割すれば, グラフは森になり, 各連結成分の頂点数を求めればいい.
色で分割したときの単純パスの総数は適当な頂点からDFSすることでで求めることができる. しかし, それだと色すべてでかかるので結局となり間に合わない.
各頂点に対してその頂点の色以外を考えるときの更新に無駄が多そう.
DFSをする際にどういう風に更新していくか.
ここで, 図1のような状況を考える. 根の色をとし, 赤のところが同じ色で塗られていると考える. 頂点の下の三角はその頂点の部分木である.
赤頂点の部分木(黒線三角)の連結成分を無視すると, 連結成分が2つあるが, 左側は図2のように求められる.
緑の三角で表された部分木のサイズ(頂点数)水色の三角で表されたすべての部分木のサイズの和
とすると, 連結成分の頂点数を求めることができることがわかる.
緑色の三角形は根の子の部分木のサイズであるのですぐ求められるが, 水色の三角形はどのように求めるか.
DFSで上に戻るとき, 色ごとに頂点を記録することが考えられるが, これだと上書きが面倒だったり頂点が多くなりすぎたりして難しい.
本来必要なものは水色の部分木のサイズの和のみであるのでそれを記録していくことを考える. こうすると図3の矢印のように根にDFSで根に帰ってきたときは色の頂点の部分木(図2における水色の三角形)のサイズの和が求められているので, 色による分割後の連結成分のサイズを求めることができる.
DFSで各色の頂点数を親に返しているわけではなく, グローバルにおいてる値(色におけるこの値をとする)を更新していってるので, 各色について頂点数を加算して, それを色の頂点数としてしまうと, 図4のような状態を考えたとき, 赤色の頂点3つ分の部分木の頂点数を加算してしまっていて, 右側の連結成分が求められない.
ここで, 根の右側の子へDFSするときにの値を記録しておく(この値をとする). こうすることで根に帰ってきたとき, 現在のとの差を見ることで頂点数の和を求めることができるようになる.
また, 図5のようにDFSで親に戻るとき部分木の頂点数をに足す必要があるが, 部分木内の他の色の部分木の頂点数をすでに足しているので余分に足してしまうことになる.
このときには, 親からDFSで来るときにの値を記録しておくことで, 親に帰るとき
(記録した値) (部分木のサイズ)
と更新すればよくなる.
また, 元の木の根で最後に全色について更新する必要があるので注意.
提出プログラム
https://atcoder.jp/contests/abc163/submissions/12185855
感想
色で分割するということは分かったが計算量を落とすことができなかった.