Codeforces Round #540 Div.3 F1: Tree Cutting (Easy Version)
問題概要
頂点からなる木が与えられ、各頂点は赤・青のどちらかで塗られているか、何も塗られていないかのいずれかである。
木の辺を 1 つだけ取り除き、同じ連結成分内には色が高々 1 種類まで登場するようにしたい。これを達成できる辺の取り除き方は何通り存在するか？
制約
木には赤の頂点と青の頂点がそれぞれ1つ以上存在する
解法
まず、木に存在する全ての赤を含む最小の連結成分と、木に存在する全ての青を含む最小の連結成分を作り、作った連結成分の頂点は赤や青に変えてしまいます。これは帰りがけdfsで、自分の子孫に赤/青の頂点があるか？を調べながら色ぬりをしてあげるとよいです。
このとき、赤の連結成分と青の連結成分の両方に含まれる頂点が存在した場合、答えとなるような辺は存在しないので、この時点で 0 を出力してプログラムを終了します。
赤と青の連結成分がそれぞれ 1 つになったあと、答えとなる辺を探索します。
赤の連結成分と青の連結成分を繋ぐパスは必ず 1 つだけ存在し、そのパスに含まれる辺が問題の条件をみたす辺になるので、パスを探してその長さを出力します。
まず、与えられる辺を順に見ていきます。
辺に繋がっている2つの頂点が同じ色の場合は、その辺は求めるパスに含まれていないことが明らかなので何もしません。2つの頂点がどちらも塗られていない場合も何もしません（この辺が求めるパスに含まれているケースも存在しますが、そのときは求めるパスの長さが 3 以上であり、この辺で探索しなくても良く、別の辺のときに求めるパスであるかを判定します）。
そうでない場合で、2つの頂点が赤と青の場合は、この辺がパスになるので 1 が答えになります。
上のどちらでもない場合は、この辺が求めるパスに含まれているかを探索します。片方の頂点が赤でもう片方の頂点は塗られていない頂点の場合、色が塗られていない頂点から赤の頂点とは反対方向に、青の頂点が見つかるまで探索します。ここで、青の頂点が見つかれば、最初の赤の頂点から見つけた青の頂点までのパスが求めるパスなので、そのパスの長さを出力します。
もし見つからなければ、この辺は求めるパスに含まれる辺ではないことがわかるので、別の辺を探索します。
「今見ている辺が求めるパスに含まれているか？」を調べる回数は、与えられる木によっては多くなることがありますが、全体でそれぞれの辺は高々2回ずつしか見ないので、辺に繋がっている2つの頂点の色が違うときに毎回調べても になります。
