ARMERIA

Rubyと競技プログラミングの話 AtCoderやCodeforcesの問題解説記事が多め。

AtCoder Beginner Contest 163 F - path pass i(マージテク解法)

F - path pass i

公式解説とは違う方法で、計算量的には  O(N(\log N)^{2}) であまり良くないのですが、私が解いた解法をまとめておきます。

解説

「通らないもの」を数える

頂点数が  N であるとき、単純パスの総数は  \frac{N(N+1)}{2} 個です。色  c について解く時には、この総数から「色  c の頂点を通らないもの」の個数を引くことにします。

 c の頂点を除いて木をいくつかの部分木に分解すると、異なる部分木間のパスは必ず色  c の頂点を含みます。そして各部分木の内部で、頂点数  n とすると色  c の頂点を通らない単純パスが  \frac{n(n+1)}{2} 個できます。これを合計すれば良いので、つまり各色についてその色の頂点で切ってできる各部分木のサイズが分かれば良いです。

今回はこれを全ての色について計算する必要があり、色ごとに木DPなどをしていては間に合いません。なので全ての色についての計算を一気にしてしまいます。

全色一気に処理するDP

適当に根を決めて木DPを行い、以下の2つの値を管理することにします。

  •  cnt\lbrack i \rbrack = 頂点  i 以下の部分木の頂点数。
  •  dp\lbrack i \rbrack\lbrack c \rbrack = 頂点  i 以下の部分木において、頂点  i の親方向から見て色  c で塞がれているような頂点の個数

後者がちょっと分かりにくいので図で具体例を示します。

f:id:betrue12:20200420001823p:plain

これを用いて部分木のサイズを集めていきます。ポイントとしては、頂点  i について見る時には色  c_{i} で区切られた部分木だけが閉じられるということです。具体的には、子  j の方向から伸びてきて頂点  i で閉じられるような部分木の頂点数は  cnt\lbrack j \rbrack - dp\lbrack j \rbrack\lbrack c_{i} \rbrack と計算できます。これを集めていけば良いです。

f:id:betrue12:20200420002354p:plain

ただしこのように「上から閉じて数える」方法だと根を含む部分木は集計されないので、最後に全ての色について根を含む部分木の個数を計算します。

この  dp\lbrack i \rbrack\lbrack c \rbrack を計算する遷移を考えます。これは子のものを全て  dp\lbrack i \rbrack に足していって、 dp\lbrack i \rbrack\lbrack c_{i} \rbrack cnt\lbrack i \rbrack で上書きすればOKです。

この  dp\lbrack i \rbrack\lbrack c \rbrack を全ての頂点×全ての色について計算しようとすると間に合いません。しかし頂点  i 以下の部分木に含まれている色の数は  cnt\lbrack i \rbrack 種類以下なので、 dp\lbrack i \rbrackmap 等で持ってしまいましょう。そうすると足し合わせていく処理において、俗称「データ構造をマージする一般的なテク」と呼ばれるテクニックが利用できます。長いのでマージテクと呼びます。

マージテクとは

マージテクとはデータ構造(C++で言うと vectorsetmap など)で保持したデータを「必ずサイズの小さい方から大きい方にマージするようにする」というルールでマージすることで操作全体の計算量を改善するテクニックです。

今回のケースでは、最初に  N 個の各頂点がそれぞれ「サイズ  1」のデータ構造を持っています。これを子から親にマージすることを繰り返して、最終的に根で「サイズ  N」のデータ構造を作ります。(実際には同じ色がある場合は消えますが)

このような時にマージの順序を気にせずに子から親にデータを移してしまうと、データ1個を移動する回数の合計は最大で  O(N^{2}) 回になってしまいます。根から一直線に伸びるようなグラフがその例です。

ここで「小さい方から大きい方に」のルールを守るために「子の  dp\lbrack j \rbrack から親の  dp\lbrack i \rbrack にマージするときに、 dp\lbrack j \rbrack のほうが大きければ事前にスワップする」という処理を入れることで、なんとデータ1個を移動する回数の合計が全部で  O(N\log N) 回になります。

この計算量は、「1回のマージで移動する要素の個数」ではなく「ある要素が操作全体で移動する回数の合計」に注目すると説明できます。要素が移動する際には元いたデータ構造よりサイズが大きいものに移動するので、「その要素が所属しているデータ構造のサイズ」が2倍以上になります。移動1回ごとにサイズが2倍以上になるという条件で最初のサイズ  1 から最後のサイズ  N になるので、その移動回数は  O(\log N) になります。これが全ての要素について言えるので、全体での移動回数合計が  O(N\log N) になる、という仕組みです。

ここで「データ構造」「データ」等の書き方をしていますが、最初に書いたように vectorsetmap やその他色々なデータ構造でこのテクニックは利用可能です(「一般的な」と呼ばれている理由です)。1つずつ要素を取り出して移し替えていく、という操作ができるものなら何でもできます。ただしその移動1回の計算量が移動回数合計と掛け算されるので、例えば map を用いると操作全体の計算量としては  O(N(\log N)^{2}) になります。

まとめ

これで全ての色について、その色の頂点で切ったときの各部分木のサイズが求められ、冒頭に書いた計算式で答えを計算できます。

ここで各色の部分木の総数が異常に増えてしまわないかというのも気になりますが、大丈夫です。ある頂点を取り除いたときに部分木の個数はその次数-1だけ増えるので、部分木の個数を全ての色について足した総数も全頂点の次数の合計、つまり  O(N) で抑えられます。

ACコード

Submission #12130625 - AtCoder Beginner Contest 163