Codeforces Round #540 F2. Tree Cutting (Hard Version)

コンテスト単位の記事がなかなか書けずにこどふぉの記事をサボってしまっているので、比較的楽に書ける1問単位の記事も書いていこうと思います。ちょっと方向性模索中。

解法

同じ色の頂点を繋ぐ

もし直接辺で繋がっていない頂点同士が同じ色だった場合、その間のパスに含まれる頂点は全て同じ連結成分に含まれる必要があります。なので前もってパス上の頂点をその色で塗りつぶしておくことで、今後の処理が楽になります。

これを効率的に行うためにLCAを利用します。同じ色の頂点全てのLCAを求めて、各頂点からLCAまでの経路上の頂点を全て塗れば良いです。（同じところを何度も塗りなおすような実装だと時間がかかるので注意）

このとき、もし塗ろうとした頂点に既に別の色が塗られていた場合は問題の条件を満たすことは絶対に不可能なので、0 を出力して終了します。ある色の頂点たちが別の色の頂点で分断されているような場合や、パス同士が交差してしまうような場合が該当します。

以降は、こうやって色を塗り足して同じ色の頂点をそれぞれ連結にした状態を考えていきます。

辺の切り方をDPで数える

適当に根を決めて木DPをします。このとき状態は以下のように取ることが出来ます。

$dp\lbrack i \rbrack\lbrack 0 \rbrack =$ 頂点 $i$ 以下の部分木における辺の切り方であって、頂点 $i$ が含まれる連結成分に有色の頂点が含まれていないような場合の数。
$dp\lbrack i \rbrack\lbrack 1 \rbrack =$ 頂点 $i$ 以下の部分木における辺の切り方であって、頂点 $i$ が含まれる連結成分に有色の頂点が含まれているような場合の数。

「辺をいくつ切ったか？」「具体的に何色の頂点が含まれているか？」というパラメータが必要そうに見えますが、それらを含めてしまうと計算が間に合いません。上手くDPをすることでこれらのパラメータを持たなくても正しい数え上げをすることができます。

先ほど同じ色の頂点は連結になるように間を塗ったので、有色の頂点の間に無色の頂点が挟まっている場合、それら有色の頂点同士は必ず異なる色になっています。そのためどの色が伸びてきているのかに関わらず、その間のどこかでは必ず辺を切らないといけません。これにより「具体的に何色の頂点が含まれているか？」という情報を不要にすることができます。

そして必要以上に辺を切らない、具体的には「 $dp\lbrack i \rbrack\lbrack 0 \rbrack$ から遷移するときには辺を切ることを考えない」ことにより、最終的には必ずちょうど $K-1$ 本の辺が切られている状態になっているはずです。これにより「辺をいくつ切ったか？」という情報を不要にすることができます。

この考え方をもとに遷移を構成していきましょう。

頂点 $i$ が有色であるとき

有色の頂点については必ず $dp\lbrack i \rbrack\lbrack 0 \rbrack = 0$ になるので、 $dp\lbrack i \rbrack\lbrack 1\rbrack$ の計算だけを考えます。それぞれの子からの遷移は、子との色関係によって

子 $j$ が $i$ と同じ色を持っている場合、 $dp\lbrack j \rbrack\lbrack 0\rbrack = 0$ であり、 $dp\lbrack j \rbrack\lbrack 1\rbrack$ から「辺を繋ぐ」という遷移があり得る。
子 $j$ が $i$ と違う色を持っている場合、 $dp\lbrack j \rbrack\lbrack 0\rbrack = 0$ であり、 $dp\lbrack j \rbrack\lbrack 1\rbrack$ から「辺を切る」という遷移があり得る。
子 $j$ が無色である場合、 $dp\lbrack j \rbrack\lbrack 0\rbrack$ から「辺を繋ぐ」という遷移と、 $dp\lbrack j \rbrack\lbrack 1\rbrack$ から「辺を切る」という遷移があり得る。

の3パターンに分かれます。これらは結局どれも $dp\lbrack j \rbrack\lbrack 0\rbrack + dp\lbrack j \rbrack\lbrack 1\rbrack$ 通りになるので、これを全ての子について掛け算したものが $dp\lbrack i \rbrack\lbrack 1\rbrack$ です。

頂点 $i$ が無色であるとき

まずは $dp\lbrack i \rbrack\lbrack 0 \rbrack$ を考えます。これは有色の頂点が含まれている連結成分を一切繋がないようにするので、

子 $j$ が有色である場合、 $dp\lbrack j \rbrack\lbrack 0\rbrack = 0$ であり、 $dp\lbrack j \rbrack\lbrack 1\rbrack$ から「辺を切る」という遷移があり得る。
子 $j$ が無色である場合、 $dp\lbrack j \rbrack\lbrack 0\rbrack$ から「辺を繋ぐ」という遷移と、 $dp\lbrack j \rbrack\lbrack 1\rbrack$ から「辺を切る」という遷移があり得る。

と、先ほどとほぼ同じになります。 $dp\lbrack j \rbrack\lbrack 0\rbrack + dp\lbrack j \rbrack\lbrack 1\rbrack$ を全ての子について掛け算したものが $dp\lbrack i \rbrack\lbrack 0\rbrack$ です。

一方 $dp\lbrack i \rbrack\lbrack 1 \rbrack$ のほうでは、どれかちょうど1つの子からは有色の頂点を含む連結成分を繋ぎ、残りの子については色を持ち込まないという遷移を考えます。この1つの子を全通り試すので、そのような子を $k$ として

$dp\lbrack i \rbrack\lbrack 1 \rbrack = \sum_{k}(dp\lbrack k \rbrack\lbrack 1\rbrack \times \prod_{j \ne k}(dp\lbrack j \rbrack\lbrack 0 \rbrack + dp\lbrack j \rbrack\lbrack 1 \rbrack))$

が遷移の計算式になります。外側の $\sum$ が全ての子 $k$ について回り、その中の $\prod$ は $k$ 以外の子 $j$ について回るので、これは単純に計算すると子の数の2乗オーダーの計算時間がかかります。ですが積の部分を $k$ を挟んで左右からの累積積を用いて計算することで1乗オーダーに落ちます。