はじめに、問題自体の解説を読みたい方は、ページ下方の「この問題の解説」から読むことを推奨します。

はじめに、この問題を解くのに必要な 重軽分解 と呼ばれるテクニックを説明します。

重軽分解

重軽分解(Heavy Light Decomposition, 以下 HLD と表記) とは、根付き木を いくつかのパスに分割するテクニックのことを言います。
HLD を利用して木をいくつかのパスに変換することで、列に対して適用できるアルゴリズムを木に対しても適用することができるようになります。

アルゴリズム

• 注：以下に説明するアルゴリズムは 正確には HLD ではなく Centroid Path Decomposition と呼ばれるアルゴリズムだと考えられますが、日本の競技プログラミング界隈ではこのアルゴリズムが HLD として広く知られており、実装も簡単なのでこちらを紹介します。

頂点 \(1\) を根とする \(N\) 頂点の根付き木を考えます。
まず、すべての頂点に対して部分木の大きさを前計算します。そして、各頂点 \(v\) について、\(v\) から子へ伸びる辺を heavy edge / light edge のいずれかに分類します。heavy edge / light edge は以下のような方法で分類します。

• 頂点 \(v\) の子のうち、部分木が一番大きい頂点を選び、\(v\) と heavy edge で結ぶ。(複数ある場合はどれでも良い)
• それ以外の子は \(v\) と light edge で結ぶ。

以下では heavy edge / light edge で結ばれた子を heavy child, light child と表現します。
このとき、各頂点から子の方向へ伸びる heavy edge は高々 1 本なので、heavy path で結ばれた頂点の集合は path graph をなします。これを heavy path と呼びます。
以上の操作により、木をいくつかの heavy path と light edge の集まりに分解することができました。

• 分解の例 (赤い太線が heavy edge, 青い点線が light edge)

性質

HLD は次の性質を持ちます。

• 任意の頂点 \(v\) から根 \(1\) までのパスは、 \(\mathrm{O}(\log n)\) 本の heavy path と light edge に分解できる。

正当性を確認します。\(v\) から根に向かって移動することを考えます。今、 light edge を経由して頂点 \(s\) から頂点 \(t\) へ移動したとします。このとき

• (\(s\) の部分木のサイズ) \(\geq\) \(2 \times\) (\(t\) の部分木のサイズ)

が成り立ちます。 なぜならば \(s\) は、\(t\) の部分木のサイズよりも大きい部分木をなす子を heavy child として持つからです。よって light edge は高々 \(\lceil \log_2 N \rceil\) 本 しか通らず、heavy edge が path graph をなすことと合わせると正当性が確認できます。

同様に、任意の頂点 \(v\) から根への path に含まれる頂点の列は、 \(\mathrm{O}(\log N)\) 本の heavy path 上の頂点の列で表すことができます。さらに LCA と組み合わせると、任意の 2 頂点 \(u-v\) 間のパスを \(\mathrm{O}(\log N)\) 本の heavy path に分解することができます。

よって、あらかじめ heavy path を列として管理しておくことで、任意のパスに対する処理を \(\mathrm{O}(\log N)\) 本の列への処理に変換することができます。

行きがけ順を利用した実装

実装は行きがけ順と併用するのが一番楽だと思います。

はじめに各頂点の部分木のサイズを計算して、 heavy edge と light edge を分類します。以下のように隣接リスト g の g[c][0] に \(c\) のheavy child を格納するのが簡潔です。

vector<vector<int>> g; // 根付き木の隣接リスト
int sub[N_MAX + 3];    // i を根とする部分木のサイズ

int dfs_hld(int c) {
  sub[c] = 1;
  for(auto& d : g[c]) {
    sub[c] += dfs_hld(d);
    // 今見ている部分木が g[c][0] を根とする部分木より大きい場合 swap
    if(sub[d] > sub[g[c][0]]) swap(d, g[c][0]);
  }
  return sub[c];
}

その後、DFS で g の行きがけ順を求めます。また、行きがけ順を計算するついでに (\(n\) を含む heavy path の根の側の端点) も計算しておきます。(以下のコードの head です)

int order[N_MAX + 3]; // i が行きがけ順で何番目に来るか
int head[N_MAX + 3];  // i を含む heavy path の端点
void dfs2(int c, int &i) {
  order[c] = i++;
  for(auto& d : g[c]) {
    head[d] = (g[c][0] == d ? head[c] : d);
    dfs2(d, i);
  }
}

すると、heavy path と行きがけ順の間に次の関係が成り立ちます。

• 任意の heavy path は行きがけ順の上でも連続した列になる。

これは、すべての \(n\) について隣接リスト g[n] の先頭要素が heavy child になっていることから従います。よって、行きがけ順を利用すれば、heavy path を容易に管理することが可能になります。

また、前計算 head の情報を利用すると、任意の頂点 \(v\) から根までのパスに含まれる heavy path を列挙することができます。(前計算でもう少し工夫すると、任意のパス \(u-v\) を結ぶ heavy path の列も求められます。考えてみましょう)

よって、例えばデータを行きがけ順に並べ替えて Segment Tree などのデータ構造に載せることで、木のパス上のデータの情報の取得・更新クエリを \(\mathrm{O}(\log N)\) 個の列への取得・更新クエリに変換できます。

参考: CodeForces の記事

コンテストで HLD を利用するケースはもう 1 つ代表的なものがあり、与えられた木に対して、いわゆるマージテク (Weighted-Union Heuristic) を行う場合が挙げられます。この問題ではマージテクに類するアルゴリズムを利用します。

この問題の解説

まず、この問題は 2 乗の木 DP で \(\mathrm{O}(N^2)\) で解けます。遷移を見ると FFT を利用した高速化が利きそうですが、いわゆるムカデグラフ (centipede graph) で \(\Theta(N^2)\) の計算量が掛かるため AC するのは難しいでしょう。

以下では \(n\) を根とする部分木についての答えの母関数を\( f_n(x)\) と表します。

まず、部分問題として次の場合を考えてみましょう。

• \(N\) 頂点の木がある。\(1\) が根で、\(1-2-3-\dots-M\) がパスになっている。また、このパスに載っていない頂点 \(v\) \((M \lt v \leq N)\) について \(f_v(x)\) が求まっている。

この問題は \(\mathrm{O}(N \log^2 N)\) で解けます。
簡潔に説明します。 (容易ではないですが、log 2 個系の練習として適切な難易度なのでわからない方は考えてみましょう)
\(g_n(x)\) を次のように置きます。

\[g_n(x) = \prod_{P_v = n, v \gt M} f_v(x)\]

すると、適切な考察により \(f_1(x)\) と \(g_{\ast}(x)\) の間に次の関係式が成り立ち、この式は分割統治を適用できます。

\[f_1(x) = \sum_{i=1}^{M+1}\left( (\text{1 if }i = M+1\text{ else }x) \prod_{j=1}^{i-1} g_j(x) \right)\]

さて、一般の木の場合を考えます。 まず、部分木の大きさを計算して heavy edge / light edge を分類します。その後、 DFS によって葉から順に DP を計算していくことを考えます。

• 余談ですが、この問題では DFS の代わりに \(n=N, N-1, \dots, 2, 1\) の順番に走査する方が言語によっては実装が容易です。(\(P_i \lt i\) より上に書いた順はトポロジカル順序の逆順なのが保証されているため)

頂点 \(n\) の heavy child を \(\text{heavy}_n\) とおきます。各頂点 \(n\) について、部分問題の “パス” を “heavy path” に置き換えた場合と同様に、 \(g_n(x)\) を \(n\) の light child の \(f_{\ast}(x)\) の総積として定義します。すなわち次の式です。

\[g_n(x) = \prod_{P_v = n, v \neq \text{heavy}_n} f_v(x)\]

頂点 \(n\) について、\(n\) を端点として子の方向へ連なる heavy path 上の頂点列を \(X_{n, 1}, X_{n, 2}, \dots, X_{n,k}\) とします。そして 、頂点列の各頂点に対する \(g_{\ast}(x)\) を並べた列 \(F_n = (g_{X_{n,1}}, \dots, g_{X_{n,k}})\) を考えます。\(f_n(x)\) のかわりにこの \(F_n\) を DP の過程で保持します。

まず、子がいない場合は \(F_n = (1)\) です。 以下では子がいる場合を考えます。
まず、あらかじめ計算してある \(F_{\text{heavy}_n}\) を (ポインタ渡しや std::move を用いて) \(\mathrm{O}(1)\) で取得して、\(F_n \gets F_{\text{heavy}_n}\) と更新します。

その後、\(n\) の light child の DP の結果をマージして \(g_n(x)\) を求めます。 \(g_n(x)\) の計算に必要な \(f_v(x)\) は、\(F_v\) を部分問題と同様の分割統治で \(f_v(x)\) に変換することで求められます。最後、\(F_n\) に \(g_n(x)\) を追加します。
この手順を DFS で葉から順に計算することで \(F_1\) を求められます。

計算量を考えます。次の 2 つの部分以外は \(\mathrm{O}(N)\) で処理できるので、これらの計算量を考えます。

• \(F_n\) を \(f_n(x)\) に変換する部分
• light edge の結果をマージして \(g_n(x)\) を計算する部分

\(F_n\) を \(f_n(x)\) に変換する部分は、\(n\) を根とする部分木のサイズを \(S_n\) として \(\mathrm{O}(S_n \log^2 S_n)\) かかります。また、\(F_n\) を \(f_n(x)\) に変換する必要があるのは light child および頂点 \(1\) のみです。ここで HLD の性質を利用すると、次の式が成り立ちます。

\[\sum_{\text{heavy}_{P_v} \neq v} S_v = \mathrm{O}(N \log N)\]

よって、この部分は DP 全体で \(\mathrm{O}(N \log^3 N)\) の計算量で抑えられるのが確認できます。また、light edge の結果をマージする部分も同様に \(\mathrm{O}(N \log^3 N)\) で抑えられます。よってアルゴリズム全体の計算量もまた \(\mathrm{O}(N \log^3 N)\) となります。

定数倍を考えると、計算量の \(\log N\) のうち 1 個は light edge の性質を利用したもので、1 個は heavy path の長さの対数を取ったもので、この 2 つの値は同時に上限に近い値を取ることはできないので、定数倍はその分軽くなります。このような理由から、このアルゴリズムは十分高速に動作すると考えられます。

実装例(C++, 454 ms)

なお、上の実装例では、heavy edge を経由したデータの受け渡しに c++ の std::move という機能を利用しています。 (C++11 以降では、関数の返り値は特定の条件を満たすと自動的に move されます。詳しくは cppreference の return statement の項 の Automatic move from local variables and parameters の章を参照してください。)

追記：より計算量の良い解法を教えてもらったので少し追記します。
light edge の結果をマージする部分は、 deque を使わずに「次数の小さい多項式を 2 個選んで積に置き換える」という方法でマージすると計算量が \(\mathrm{O}(N \log ^2 N)\) で抑えられるのが示せるようです。また、heavy path の値を分割統治で求める部分も、分割の仕方を上手く決めることで \(\log\) を 1 個落とせて、結局 全体で \(\mathrm{O}(N \log^2 N)\) で計算できるようです。(細かい部分はわかっていません)

• 参考：Nachia さんの記事 　yosupo さんの記事