この問題はいくつかの解法がありますが、この解説では Static top tree と呼ばれるデータ構造を利用した \(\mathrm{O}( (N+Q) \log N)\) 解法を説明します。

目標 : 木 DP をセグ木に載せる !?

はじめに、この問題を解くためにはどのような操作が必要となるかを考えてみます。

まず、この問題をクエリあたり \(\mathrm{O}(N)\) の計算量で解くことを考えてみましょう。これは簡単で、典型的な木 DP を行えばよいです。すなわち、次のような \(f(n)\) をボトムアップに計算するコードを書いて、\(A\) を更新するたびに毎回 (頂点 \(0\) を根として) calc_dp(0) を呼び出せばよいです。しかしながらこの方法では \(\mathrm{O}(NQ)\) の計算量がかかり TLE してしまいます。

using mint = atcoder::modint998244353;

vector<vector<int>> g; // 根付き木の隣接リスト
mint A[n_max];
mint calc_dp(int v) {
  if(g[v].empty()) return A[v];
  mint prod = 1;
  for(auto& child : g[v]) prod *= calc_dp(child);
  return A[v] + prod;
}

さて、この問題は次の 2 種類のクエリを処理する問題とみなすことが出来ます。

問題をクエリ形式で言い換えたもの

• \(N\) 頂点の根付き木と数列 \((A_1, A_2, \dots, A_N)\) が与えられる。次の 2 種類のクエリを処理
  • \(A_v\) の値を \(x\) に更新
  • 上記の木 DP を計算して、\(f(1)\) を出力

つまり、この問題は木 DP に関して「頂点の値を更新」「全体の値を取得」という 2 種類のクエリを要求しています。
このような形式の問題は木においては珍しいですが、列に関する問題では頻出なのは皆さんもご存じの通りかと思います。

典型的な列の問題

• 列 \(A_1, A_2, \dots, A_N\) が与えられる。次の 2 種類のクエリを処理
  • \(A_i\) の値を \(x\) に更新
  • \(A_1 \oplus A_2 \oplus \dots \oplus A_N\) を計算 (\(\oplus\) は何らかのモノイド和)

列の問題は、セグメント木を利用することでクエリあたり \(\mathrm{O}(\log N)\) で高速に処理することができるようになります。また、木の場合でもクエリが列に関するクエリに帰着できる場合は (例 : パス上の頂点に載っているモノイドの和) 、重軽分解とセグメント木を組み合わせて処理することが出来ます。

一方、元の問題は、木 DP を計算する過程をデータ構造とみなした時に、木 DP の 1 点更新・全体取得に相当する操作が要求されています。この問題はいわば、「木 DP をセグメント木に載せる」ような処理が求められているとも言えるでしょう。(当然ですが通常のセグメント木には載りません)

では、木 DP をどのようにしてセグメント木状のデータ構造に載せればよいかを考えてみることにします。

考察 1 : 完全二分木の場合の解法

簡単なコーナーケースとして、根付き木が完全二分木の場合に問題を解く方法を考えてみましょう。

上述した木 DP を付け焼き刃的に高速化する方法として次のようなコードが考えられます。

using mint = atcoder::modint998244353;

vector<vector<int>> g; // 根付き木の隣接リスト
vector<int> P;         // 親の頂点 (ただし P[0] = -1)
mint A[n_max], f[n_max];
mint precalc(int v) {
  if(g[v].empty()) return f[v] = A[v];
  mint prod = 1;
  for(auto& child : g[v]) prod *= precalc(child);
  return f[v] = A[v] + prod;
}

void update(int v, mint x) {
  A[v] = x;
  for(; v != -1; v = P[v]) {
    if(g[v].empty()) f[v] = A[v];
    mint prod = 1;
    for(auto& child : g[v]) prod *= f[child];
    f[v] = A[v] + prod;
  }
}

コードの概要を説明します。 まず、あらかじめ前計算で \(f(1), f(2), \dots f(n)\) の値をメモしておきます。そして、A[v] の値を更新したときに \(f(n)\) の値が変わるのは \(v\) および \(v\) の祖先の頂点だけなので、それ以外の頂点に関する情報は前計算したものを使用しながら、\(v\) の祖先の情報をボトムアップに再計算するというものです。

この解法は特定のケースに対しては非常に高速に動作するようになります。それが根付き木が完全二分木である場合です。完全二分木は深さが \(\mathrm{O}(\log N)\) であり、各頂点の子が高々 \(2\) 個しかなく各頂点の再計算も \(\mathrm{O}(1)\) で済むため、クエリあたり \(\mathrm{O}(\log N)\) で DP の結果を再計算することが出来ます。

もちろん、一般の木ではこの解法は最悪計算量が \(\mathrm{O}(NQ)\) のままで改善できていません。具体的には、

• 深さが最悪 \(\mathrm{O}(N)\) である
• \(1\) つの頂点が持つ子の個数が最悪 \(\mathrm{O}(N)\) である

という 2 つの要因により最悪計算量がクエリあたり \(\mathrm{O}(N)\) のままです。この解法を応用するためには、一般の木の持つ上記の 2 つの性質をクリアするような、何らかの深さが \(\mathrm{O}(\log N)\) の二分木に変換する必要がありそうです。

考察 2 : 木の分解と木 DP の関係性

考察を発展させるために、木 DP を別の角度から解釈してみます。実は、木 DP は実は「部分木をマージして根付き木を生成する過程に情報を載せたもの」と捉えることが出来ます。今からこのことを説明します。

マージにより木を生成する過程を考えるには、その逆、木を分解する過程を考えるのが有用です。そこではじめに木を分解する過程を説明します。
次の (1) ～ (3) を木から辺が無くなるまで再帰的に繰り返すことで、根付き木を頂点と辺に分解することができます。(下図も参照してください)

• (1) 根の頂点を取り除く。ただし根に隣接する辺は取り除かず、そのような辺は「情報を持たない頂点」を端点の一方に持つものとみなす。 (このような頂点を以降では virtual な頂点 と呼ぶ。)
• (2) virtual な頂点で集約されている部分木たちを、virtual な頂点を分裂させることで分割する。
• (3) 各部分木の virtual な頂点および virtual な頂点に隣接する辺を取り除く。部分木は再び通常の根付き木になる。

この操作手順を逆順に行っていくことで、部分木がだんだんとマージされていき最終的に 1 つの根付き木を得ることが出来ます。
マージする過程を関数として表してみましょう。次の 4 種類の関数を定義します。以降では、頂点は情報を持ち区別できるが、辺同士は情報を持たず区別できないものとします。

• vertex(v) : 頂点 \(v\) を生成する関数。
• add_vertex(t, v) : (1) の逆手順を行う関数。つまり、根が virtual な木 \(t\) の根に \(v\) を代入する関数。
• merge(x, y) : (2) の逆手順を行う関数。つまり、virtual な根を持つ 2 つの木 \(x, y\) を 1 つにする関数。
• add_edge(t) : (3) の逆手順を行う関数。つまり、根付き木 \(t\) に virtual な根を生やす関数。

そして、関数 generate_tree(v) を、隣接リスト \(g\) が表す根付き木の頂点 \(v\) を根とする部分木を生成する関数とします。すると、generate_tree(v) は次のようになります。

vector<vector<int>> g; // 根付き木の隣接リスト

Tree generate_tree(int v) {
  if(g[v].empty()) return vertex(v);
  vector<Tree> children;
  for(auto& child : g[v]) {
    Tree t = generate_tree(child);
    children.push_back(add_edge(t));
  }
  Tree t = children[0];
  for(int i = 1; i < (int)children.size(); i++) {
    t = merge(t, children[i]);
  }
  return add_vertex(t, v);
}

この部分木をマージして新たな部分木を生成する手順は木 DP の手順と一致しています。例えば先に上げた木 DP は、抽象化することで次のように書き換えられます。

using mint = atcoder::modint998244353;

using T = mint;
T vertex(int v) { return A[v]; }
T add_vertex(T x, int v) { return x + A[v]; }
T merge(T x, T y) { return x * y; }
T add_edge(T x) { return x; } 

vector<vector<int>> g; // 根付き木の隣接リスト

T calc_dp(int v) {
  if(g[v].empty()) return vertex(v);
  vector<T> children;
  for(auto& child : g[v]) {
    T t = calc_dp(child);
    children.push_back(add_edge(t));
  }
  T t = children[0];
  for(int i = 1; i < (int)children.size(); i++) {
    t = merge(t, children[i]);
  }
  return add_vertex(t, v);
}

generate_tree(v) 関数と calc_dp(v) 関数の 2 つが全く同じ形をしていることが注視すれば確認できると思います。

このような視点から考察すると、木 DP は部分木のマージ過程に情報を載せたものとして解釈できます。詳しく説明すると次のようになります。

• 葉の頂点からスタートして「頂点の追加」「辺の追加」「根が virtual な根付き木同士のマージ」という 3 種類の操作を繰り返して部分木をマージしていくことを考える。
• マージする過程で生成される部分木全てに対して、部分木に関連する何らかの情報を定義する。
• そして、情報同士をマージする関数を適切に定義することで、部分木のマージによって新たに生成される部分木に対応する情報を計算することができる。

さて、今までの話をまとめると次のようになります。

• 木が深さ \(\mathrm{O}(\log N)\) の二分木であれば\(\mathrm{O}(\log N)\) で木 DP を再計算できる
• 一般の木は次の 2 つの特徴を持ち、これが効率的な再計算を阻害する
  • 深さが最悪 \(\mathrm{O}(N)\) である
  • \(1\) つの頂点が持つ子の個数が最悪 \(\mathrm{O}(N)\) である
• ところで木 DP は部分木をマージする過程に情報を載せたものである
  • 部分木をマージする過程は、木を分解する過程の逆を辿ることで生成できる

この 2 つの事実から次の帰結を得ることが出来ます。

• 深さが \(\mathrm{O}(\log N)\) で二分木状である部分木のマージ過程を生成できれば、それを辿ることで効率よく木 DP を再計算できる？

このような木のマージ手順を実現したものが Static top tree になります。

Static top tree

注意 : Static top tree は Top tree の名を冠していますが、(狭義の) Top tree とは全く構造が異なるデータ構造です。Top tree を勉強したい方は混同しないようご注意ください。

• より正確には、Static top tree は「Link/Cut tree の normal edge が持つ情報を Splay tree でまとめ上げることで出来る、Top tree の機能のほとんどを代替できるデータ構造」(広義の Top tree) を static にしたものです。
  • なお、(狭義の) Top tree を static にすることで Static top tree を構成することも可能で、そのようなデータ構造を Static top tree と呼ぶ人もいます。ここでは説明を割愛しますが、興味がある方は tatyam さんの実装およびコメント を参考にしてください。

Static top tree とは、「部分木をマージする過程を表した深さ \(\mathrm{O}(\log N)\) の二分木」です。(正確には「部分木」ではない。後述)

Static top tree を構成する方法を説明するために、まずはマージする過程の逆、木を分解する手順を説明します。
はじめに木を重軽分解して各辺に heavy edge, light edge を割り当てます。(重軽分解を知らない方は ABC269 Ex の解説 を参考にしてください。)
そして、次の (1) ～ (4) を木から辺が無くなるまで再帰的に繰り返します。(先に挙げた分解の手順と異なる部分を 太字 で表しています。図も参照してください)

• (1) 根につながる heavy path を選び、そこから heavy edge を取り除く。
• (2) 根の頂点を取り除く。ただし根に隣接する辺は取り除かず、そのような辺は virtual な頂点を端点の一方に持つものとみなす。
• (3) virtual な頂点で集約されている部分木たちを、virtual な頂点を分裂させることで分割する。
• (4) 各部分木の virtual な頂点および virtual な頂点に隣接する light edge を取り除く。部分木は再び通常の根付き木になる。

ここで注記しておく点として、分解の過程で登場するグラフは根付き木の部分木に限らないという点が挙げられます。というのも、(1) の手順で辺を取り除く時に、辺を取り除く順番によっては途中で 「根付き木と呼べないもの」が発生するからです。そのため、以降では Top tree の用語を借りて、分解の過程で得られるグラフを cluster と呼ぶことにします。

図を見ると分かる通り、木を分解する過程において次の 2 種類の cluster が生成されます。

• virtual でない頂点を根とする部分木が \(0\) 本以上の heavy edge でパス状につながってできる cluster
• virtual な頂点を根とする部分木の形をした cluster

こちらも Top tree の用語を借りて、前者を path cluster 、後者を point cluster と呼ぶことにします。(根につながる heavy edge が存在しない path cluster は path とは少し呼びづらいですが、ここでは気にしないことにします)

次に、分解する手順を逆順に辿ることで cluster を二分木状にマージしていくことを考えます。 分解する過程において、cluster 同士の merge は「path cluster 同士の merge」「point cluster 同士の merge 」の 2 種類が発生します。これもまた Top tree の用語を借りて、path cluster 同士の merge を compress 、point cluster 同士の merge を rake と呼ぶことにします。

さて、cluster を二分木状にマージする際に一工夫を入れることで深さを \(\mathrm{O}(\log N)\) に抑えるのが Static top tree のポイントです。とはいえ一工夫入れると言っても (2), (4) の逆手順は工夫の入れようがないので、工夫を入れる余地があるのは

• (1) 根につながる heavy path を選び、そこから heavy edge を取り除く。
• (3) virtual な頂点で集約されている部分木たちを、virtual な頂点を分裂させることで分割する。

の逆手順です。これらは直感的には次の戦略に従って cluster をマージしていくのが良さそうです。

• (1) の逆手順 : path cluster を compress でマージしていく時に、マージ過程が完全二分木状になるようにマージする
• (3) の逆手順 : point cluster を rake でマージしていく時に、マージ過程が完全二分木状になるようにマージする

これは非常に合理的な戦略です。なぜならば、一般の木が持つ効率的な再計算を阻害する 2 つの性質は

• 深さが最悪 \(\mathrm{O}(N)\) である
• \(1\) つの頂点が持つ子の個数が最悪 \(\mathrm{O}(N)\) である

という 2 つでしたが、前者を compress によるマージが、後者を rake によるマージが解消することになるからです。
しかしながら、この戦略は最悪ケースで木の深さが \(\mathrm{O}(\log^2 N)\) になってしまいます。(詳細は略します。重軽分解 + セグメント木でパスクエリを処理する際に自然な実装だと worst \(\mathrm{O}(\log ^2 N)\) かかるのと同じ理屈です。)

そこで、さらにもう一工夫します。cluster を完全二分木状になるようにマージするとは、言い換えると「左右の子の含む cluster の個数ができるだけ等しくなるようにマージする」という操作になります。これを少し変更して、「cluster が含む頂点数ができるだけ左の子と右の子で等しくなるようにマージする」ということにします。実は、そのようにすることでマージ過程全体を表す木の深さを \(\mathrm{O}(\log N)\) に抑えられることが証明できます。(こちらも詳細は略します。重軽分解 + セグメント木でパスクエリを処理する際に少し工夫すると worst \(\mathrm{O}(\log N)\) になるのと同じ理屈です。参考 : Nachia さんの記事, errorgorn さんの記事の “Balanced HLD” の項)

以上の方法により、cluster をマージする過程を深さ \(\mathrm{O}(\log N)\) の二分木で表すことが出来ました。

実装例(5 行目～ 74 行目) : 実装の際は maspy さんと tatyam さんの実装を参考にしました。

元の問題の解法

さて、Static top tree を使用して元の問題を解く方法を説明します。
cluster をマージする過程で必要なのは次の 5 種類の関数です。

• vertex(v) : 頂点 \(v\) のみからなる path cluster を生成する関数。
• compress(p, c) : (1) の逆手順を行う関数。つまり、path cluster \(p, c\) (\(p\) が根に近い側にある) をマージする関数。
• add_vertex(t, v) : (2) の逆手順を行う関数。つまり、point cluster \(t\) の根に頂点 \(v\) を代入して path cluster にする関数。
• rake(x, y) : (3) の逆手順を行う関数。つまり、point cluster \(x, y\) をマージする関数。
• add_edge(t) : (4) の逆手順を行う関数。つまり、path cluster \(t\) に virtual な根を生やして point cluster にする関数。

木 DP を行う時と同様に、木に載せる情報を定義して、これら 5 種類の関数に対応する DP の遷移を構成することでこの問題を解くことが出来ます。

この中で一番難しいのは path cluster 同士のマージである compress です。なぜならば、point cluster は木 DP における部分木と対応しているため、木 DP 同様の情報を載せて rake でも同様の遷移を実装すれば基本的にどうにかなる一方、path cluster 同士のマージは「変な形をした何か」同士のマージになるからです。

ここで、cluster の外部との接点を、Top tree の用語を借りて boundary vertex と呼ぶことにします。path cluster は基本的に「根に近い方の boundary vertex」「根から遠い方の boundary vertex」という 2 つの boundary vertex を持ちます。(path cluster が根付き木の場合は例外で、2 つの boundary vertex が同一の頂点になります。) 問題によっては、この 2 つの boundary vertex に注目することで DP の遷移を構成しやすくなります。

適切な観察により、path cluster は「根から遠い方の boundary vertex にハッシュ値が \(x\) である根付き木が結合した時に、根に近い方の boundary vertex を根とする根付き木のハッシュ値は \(ax+b\) になる」となる値 \((a, b)\) を持てば良いことがわかります。このように path cluster に載せる情報を定義すると compress はアフィン関数の合成として定義することが出来ます。その他の関数も適切な考察により次のような関数を実装すればよいことがわかります。

using mint = atcoder::modint998244353;

vector<int> A;
struct Path {
  mint a, b;
};
using Point = mint;
Path vertex(int i) { return {1, A[i]}; }
Path compress(Path p, Path c) { return {p.a * c.a, p.a * c.b + p.b}; };
Point rake(Point l, Point r) { return l * r; };
Point add_edge(Path d) { return d.b; };
Path add_vertex(Point d, int i) { return {d, A[i]}; };

Static top tree を適切に抽象化すれば、基本的に上の 5 種類の関数を実装するだけで問題を解くことが出来るようになります。これはセグメント木が 2 種類の関数を実装すれば使えるのに似ていますね。

• 余談ですが、セグメント木は当初は完全二分木に限らず平衡な二分木全般に対するアルゴリズムとして定義されていたようです。(参考)このような解釈に則ると、Static top tree を用いた木 DP の更新はセグメント木の類似と捉えることができそうです。

後はクエリを処理する関数を実装すればよく、これは DP を計算・更新する過程を適切に実装すればよいです。

実装例(76 行目～ 134 行目)

Static top tree の応用

Static top tree の応用例を簡単にまとめます。

この問題では初等的な DP を Static top tree に載せただけでしたが、木 DP は非常に様々な種類があります。どのような木 DP であっても、compress をうまく定義することができれば木 DP を Static top tree に載せることが出来ます。
例えば以下の 2 つの問題は元の問題とは毛色の異なる問題ですが、 どちらも Static top tree を利用すればクエリあたり \(\mathrm{O}(\log N)\) で処理することが出来ます。(練習問題として考えてみましょう)

(木における最大独立集合) : \(N\) 頂点の木がある。各頂点には B, R, ? の 3 状態のいずれかが割り当てられている。次の 2 種類のクエリを処理

• 頂点を 1 つ選んで状態を変更
• 次の条件を全て満たす頂点の集合の最大サイズを出力
  • 状態が B の頂点は集合に含まれる
  • 状態が R の頂点は集合に含まれない
  • 集合に含まれるどの 2 個の頂点も木上で隣り合っていてはならない

(木の直径) : \(N\) 頂点の辺重み付き木がある。次の 2 種類のクエリを処理

• 辺を 1 本選んで重みを変更
• 木の直径の長さ、および直径のうちどれか 1 つの両端点を取得
  • ヒント : 辺重み付き木を Static top tree に載せる場合は「Static top tree を辺の情報も扱えるように改造する」「辺を頂点とみなした \(2N-1\) 頂点のグラフを考える」という方法が考えられます(後者の方が追加実装が少なく済むため楽です)

高難度コンテストにおける Static top tree を適用可能な問題の出題例もいくつか載せておきます。

出題例(UCup, JOI, Codeforces div.1 ネタバレ注意)

• The 2nd Universal Cup. Stage 19 F: When Anton Saw This Task He Reacted With 😩
• JOI 2023/2024 春季トレーニング: JOI ツアー
• CodeTON Round 8 H: Farmer John's Favorite Intern

また、Static top tree をうまく利用すると、木 DP だけでなく全方位木 DP の 1 点更新が可能になります。具体的には次のような問題も解くことが出来ます。(詳細は略します)

SPOJ QTREE6 ：\(N\) 頂点の木がある。各頂点は黒か白で塗られている。次の 2 種類のクエリを処理

• 頂点を 1 個選んで色を反転
• 頂点を 1 個選んで、その頂点と連結な頂点の個数を出力。ここで頂点 \(u\) と頂点 \(v\) が連結であるとは、頂点 \(u\) から頂点 \(v\) までのパス上の頂点 (両端点も含む) の色が全て一致していることを言う

Static top tree の用途は木 DP の更新の高速化だけではありません。
列に対する操作と木に対する操作の対比を考えてみましょう。Static top tree は列で言うところの「列をセグメント木状に分割する操作」に相当します。これらは深さが \(\mathrm{O}(\log N)\) の二分木状に全体をマージしていく木構造のデータ構造である点が共通しています。このような対比を考えると、列をセグメント木状に分割する操作を経て出来るアルゴリズムは、木でも同じことが出来ると予想されます。例えば今までに扱った「木 DP の 1 点更新」は「セグメント木の 1 点更新・全体取得」の木 version と言えます。
列をセグメント木状に分割する操作はセグメント木以外でも有用です。例えば、列の分割統治では列をセグメント木状に分割します。逆に言うと、Static top tree 上で分割統治をすることで木上の分割統治を行うことが出来る、ということになります。具体的には、Static top tree を利用すれば次の問題を \(\mathrm{O}(N \log^2 N)\) で解くことが出来ます。

• ABC269Ex Antichain : \(N\) 頂点の根付き木がある。(頂点 \(a\) が頂点 \(b\) の子孫) \(\iff (a \leq b)\) となるように二項関係を定義する。\(k=1, 2, \dots, N\) について、サイズ \(k\) の Antichain の個数 \(\text{mod }998244353\) を計算

• yukicoder No.2595 Parsing Challenge : 1 ~ 9, +, -, *, (, ) からなる長さ \(N\) の数式を計算 (ヒント : 構文木を考える)

このような観点から考えると、 Static top tree には「木 DP を 1 点更新」だけでなく幅広い用途があると考えられます。

さて、現代の競技プログラミングでは「平衡二分木」や「Link/Cut tree」はコンテストに出題されることはほとんどありませんが、これらの機能を制限して static にしたデータ構造である「セグメント木」や「重軽分解 + セグメント木」は主要なアルゴリズムの 1 つに数えられるほどにコンテストで出題されています。
Static top tree は短時間コンテストで tier 1 のアルゴリズムになることは無いと予想していますが、セグメント木と同様に抽象化すれば軽実装になるのは魅力的で、さらに載せる演算が少し非自明なのはいくらか面白味があるので、今後何かの拍子で Static top tree が Top tree 状のデータ構造の static 版として局所的にメジャーになることもあるかもしれませんね。