yukicoder No.2116 Making Forest Hard を平方分割+rerootingで解く

お久しぶりです。最近はTTPC2022のtesterやったりICPC観戦したりして急に競プロ熱が再燃しています。

ところで今回はyukicoder No.2116 “Making Forest Hard” を平方分割+rerootingで解こうと思います。

https://yukicoder.me/problems/no/2116

問題概要

$N$ 頂点の木 $T$ がある
各頂点には整数 $A_i$ が書かれている
木 $T$ の任意の連結成分 $U$ について、 $U$ 内の頂点数を $S(U)$ 、 $U$ 内の頂点の最大の $A_i$ の値を $M(U)$ と定義し、連結成分 $U$ のスコアを $g(U) = S(U)M(U)$ と定義する
木 $T$ の辺集合 $E$ の任意の部分集合 $E'$ を考え、 $T$ から $E'$ の辺をすべて削除した木(連結成分)の集合を $T'(E')$ とする
$T'(E')$ のスコアを、任意の $T'(E')$ 内の連結成分 $U$ のスコアの総和 $f(E') = \sum_{U\in T'(E')} g(U)$ と定義する
任意の部分集合 $E'\subseteq E$ に対する $T'(E')$ のスコアの総和 $\sum_{E'\subseteq E}f(E')$ を MOD で剰余を取って出力せよ

基本考察: 寄与

最大値(×サイズ)の任意の操作における総和を考える際、典型的な方法の1つが寄与の計算です。今回はある頂点 $i$ に書かれている値 $A_i$ がどの程度スコアに寄与するかを考えます。

なお、 $M(U)$ の計算のために使う最大値計算を、頂点IDによってタイブレークしてどの頂点の値かを一意に定めることにします。つまり $M(U) = A_{\arg\max_{i \in U} (A_iN + i)}$ のようにします。

以降はこのタイブレークされた値を考え、 $A_i$ は全て相違なるものと仮定し、単純に $A_i$ の大小で記述します。

木DP O(N²)

さて、頂点 $i$ に着目するため、一度 $i$ を根として木を見、これを $R_i$ とします(元々の木 $T$ は根の決まっていない無向木です)。

例として適当な木 $R_i$ を挙げます。 $i=9$ に着目しています。

なお簡単のため、この図では頂点ID = $A_i$ としています。実際の問題ではそのような制約は無いので留意してください。

またグレーに塗られた頂点は $A_v>A_i$ となる頂点です。これらの頂点は連結成分に含まれないようにする必要があります。

この根付き木 $R_i$ において、各頂点 $v$ について

$v$ を根とした $R_i$ の部分木(連結成分) $U$ で、 $M(U) \le A_i$ であるようなものの通り数 $X_v$
そのような $U$ についての $S(U)$ の総和 $Y_v$
$v$ を根とした $R_i$ の部分木に含まれる辺の数 $Z_v$

を求める木DPを考えます。

まず辺の数 $Z_v$ はそのまま辺の数を足し上げるだけで良いです。

次に $X_v, Y_v$ について、頂点 $v$ に対して子 $w$ の部分木の結果を順番にマージした際の値の変化を考えます。子 $w$ の手前までマージしていった現在の頂点 $v$ の値を $(X'_v,Y'_v)$ とし、子 $w$ をマージした後の値を $(X''_v,Y''_v)$ とします。子 $w$ の値は $(X_w,Y_w,Z_w)$ とします。

辺を加え入れる場合
- 通り数は $v$ 側と $w$ 側の値を乗算することになり $X''_v \leftarrow X'_vX_w$ となる
- サイズの総和は $v$ 側の総和 $Y'_v$ が $X_w$ 通り寄与し、 $w$ 側の総和 $Y_w$ が $X'_v$ 通り寄与するため、 $Y''_v \leftarrow Y'_vX_w+X'_vY_w$ となる
辺を取り除く場合
- 子 $w$ の部分木内に含まれる辺はどのようにしても良くなる
- そのため $X''_v \leftarrow X'_v\mathrm{pow}(2,Z_w)$ 及び $Y''_v \leftarrow Y'_v\mathrm{pow}(2,Z_w)$ となる

以降、簡単のため $W_v := \mathrm{pow}(2,Z_v)$ と表記します。また毎回 pow を計算するのも面倒なので実装上も基本的にはこの形で持つことにします。

この2通りを足し合わせれば良いですが、 $A_w>A_i$ となる場合は辺 $(v,w)$ を取り除かなければなりません。よって

の時
- $X''_v=X'_vW_w+X'_vX_w$
- $Y''_v=Y'_vW_w+Y'_vX_w+X'_vY_w$
の時
- $X''_v=X'_vW_w$
- $Y''_v=Y'_vW_w$

となりますが、 $A_w>A_i$ の時に $X_w=Y_w=0$ とすれば上側の計算だけで済みます。またこれらの計算は全て $X'_v,Y'_v$ の線形和になるので、初期化のタイミングで $X_v=Y_v=0$ とすれば場合分けが除去できます。ちなみに $A_w\le A_i$ の場合は $X_v=Y_v=1$ と初期化します。

木DPの結果、 $Y_i$ に $A_i$ の寄与する乗数が求まるため、答えに $A_iY_i$ を足し合わせます。

以上で、各頂点 $i$ に対して $O(N)$ のDFSで寄与が求まるため、全ての頂点に対して行えばトータル $O(N^2)$ の計算量で求まりました。

発展考察: 平方分割+rerooting

さて、木DPを高速化するにあたり、典型の一つである全方位木DPを考えてみます。が、上述の木DPは根の値 $A_i$ に応じてグレーに塗られた切らなければいけないノード $A_v>A_i$ が変わるという問題点があります。

逆に $A_v>A_i$ という制約に対処しようと考えると $A_i$ の昇順もしくは降順で木DPすることで、変化する頂点は根に来るため $O(1)$ の変化でほとんど同じ木DPを解くことが出来るようになりますが、逆にrerootingのコストが1回につき最大で直径分かかります。

そこで、平方分割によって木の直径を $O(\sqrt N)$ に圧縮し、rerootingのコストを1回あたり $O(\sqrt N)$ とし、トータルコストを $O(N\sqrt N)$ にすることを目指します。

平方分割

平方分割した状態でrerootingとクエリ対象の頂点編集を行いたいので、次のような性質が必要です:

クエリする頂点集合を $D={i_1,i_2,\dots,i_B}$ とする
頂点 $i\in D$ は圧縮しない
平方分割後の木の直径が $O(\sqrt N)$ である
木DPの値の整合性を保ったまま回転が出来る

このように平方分割するには次のようにするのが良いです:

$i_1$ を根とするような木を考える
任意のについて、も圧縮しない
- このとき、圧縮しない頂点を分割点と呼び、その集合を $D'$ とする
- このとき $i_1$ が根であることに注意すると $\lvert D'\rvert \le 2B-2$
によって分割された連結成分について、それらは1個か2個の分割点と繋がっている
- 0個: 元々その連結成分が木全体である場合のみ起こり得るが $B\ge 1$ のため存在しない
- 3個以上: そのような連結成分があると仮定すると、連結成分内のある1つの頂点から相違なる3つの分割点 $i_x,i_y,i_z\in D'$ への交わらないパスが取れるが、このときその頂点は3つの分割点のどれかのLCAであり、 $D'$ に属さない頂点から構成されるという仮定に矛盾するため、存在しない
- (これはつまるところtop treeにおけるclusterとboundary vertexである)
2個の分割点と繋がっている連結成分はたかだか個
- つまり $\le 2B-3$ 個

下図はクエリする頂点集合を 9~12 としたときの木のサンプルです。

青い 9~12 の頂点が $D$ 、赤い 7 の頂点が $D' \setminus D$ 、太めの黒い曲線に囲まれた連結成分が2点の分割点に隣接する(圧縮対象の)連結成分です。

分割点では通常の木DPを計算することとし、連結成分に対する木DPを効率化します。

全ての分割された連結成分について、隣接する分割点を根方向としたときの木DPの値を計算できるようにする
- 隣接1個: rerootingによって根方向が変わることは無いため、特に何もしなくて良い(極論圧縮もしなくて良い)
- 隣接2個: rerootingによって根方向が変わることがあり、たかだか1個の分割点を部分木に持つ。この分割点以下の木DPの値を入力とし、もう一方の分割点を根方向とした連結成分の木DPの値を出力する作用素を構成する必要がある

この作用素について考えます。

木DPの圧縮

元の木DPを再掲します:

$W''_v=2W'_vW_w$
$X''_v=X'_vW_w+X'_vX_w$
$Y''_v=Y'_vW_w+Y'_vX_w+X'_vY_w$

葉方向の分割点を $u$ とすると、木の回転などによって $(W_u,X_u,Y_u)$ の値は変化するため、これらの値を自由変数と考えると、

分割点 $u$ : $(X_u,Y_u,W_u)$ という値は $(W_u,X_u,Y_u)$ の線形結合(恒等関数)で書ける
分割点 $u$ を子孫に含まない頂点 : $(W_u,X_u,Y_u)$ に依存しない定数値で書ける
分割点を子孫に含む頂点 : 分割点を含む部分木がの線形結合で書けるなら、定数とのマージなので計算結果もの線形結合で書ける
- $u$ の深さに関する帰納法で証明できる

となるため、この線形結合の係数行列さえ保存すれば木の回転によって根方向が変わった場合でも木DPの値を高速に更新することが出来ます。

木DPの計算式から上三角行列となるため、係数保持に必要な値は6個で十分です(実際には5個でも大丈夫です)。

木DPのrerooting

次に木の回転(rerooting)を考えます。

rerootingの典型として「逆演算を定義し、回転する部分木を打ち消す」という手法が考えられます。一方で今回は剰余乗算を含んでいるため、0が乗算される際の扱いが困難です。また剰余除算に $O(\log\mathrm{MOD})$ がかかってしまうのもネックです。

今回は平方分割によって主要な木の頂点数が $O(B)$ となっているため、そもそも任意のパス上の次数の総和も $O(B)$ になりそうです。

ただし分割点との隣接が1点のみの連結成分の数は抑えられていません。これらの部分木については値の更新が行われないため、木の回転のたびに計算しないようにすることで、1回の rerooting が $O(B)$ で収まります。

実装方針

根を決める
- 最初のクエリは $A_{i_1}$ について行うことに留意
$D$ の頂点を塗っておく
圧縮のDFSを行う: 各頂点について
- の子孫にの頂点を含むか、の頂点を含む子の数を求める
  - $D'$ の頂点を含む子の数が1の場合は圧縮対象、2以上の場合は分割点となる
- 通常の木DPの値 $(W_v,X_v,Y_v)$ 及び $D'$ の頂点を含まないstaticな木DPの値 $(W^s_v,X^s_v,Y^s_v)$ を求める
- 圧縮対象の分割点へのリンクと係数行列を計算する
  - 圧縮対象のパスの長さが1の場合は圧縮しないことにするとちょっと楽
について
- を根にする: の親を根にする(再帰) → が浅くなるように回転する
  - 回転は圧縮ノードを間に挟むか否かで若干変わる
- $i$ についてクエリ結果を求める
- $i$ を変化させる (今回は $A_i > A_{\mathrm{next}\ i}$ となるように変化させる等)

解答ソース

Easy : https://yukicoder.me/submissions/828995
Hard : https://yukicoder.me/submissions/828996

まとめ

ちなみにオフライン限定ですが、直径が $O(\sqrt N)$ のためパス系クエリや、圧縮をうまく取り扱えば遅延伝播系クエリ、更に森に対して圧縮すればLink Cutにも対応出来ると思います。

まぁ多分top-treeで良いんですが、しいて言えば以下のメリットを見いだせると思います:

「木を木で取り扱う」のに対し、このアルゴリズムは連結成分の圧縮しかしない
- 概念が分かりやすく、デバッグやリファクタリングも比較的楽な部類かと思います
定数倍は軽い(はず)
- オーダーはめちゃめちゃ悪い、まだ最適化の余地がありそうですが N=100000 が限度だと思います
実装も軽い(はず)
- 多分……オフラインだからその取り扱いが煩雑かも

補遺: 経緯

今回のアルゴリズムはTTPC2022のtesterをしている際に思いついたものになります。TTPC2022 testerに呼んでくださった運営の方ありがとうございました。今更ですが参加してくださった皆様もありがとうございました。

実はもう1問原案を担当していたのですが⋯ pic.twitter.com/6bipTytdYh
— ⚡すく⚡ (@0214sh7) 2022年11月5日

まぁTTPC2022の木の問題はD: XOR Tree Pathだけだったんですけどね。

参考文献

木の平方分割、出来るもの出来ないもの - noshi91のメモ
- https://noshi91.hatenablog.com/entry/2019/12/07/140433
- 今回使った平方分割は例1の亜種と言えそうです
top tree 概要 - noshi91のメモ
- https://noshi91.hatenablog.com/entry/2019/08/05/175545
- 概念を考えるにあたって参考になりました
- 列と木の対比と言えば、列のマージ・分割に対するオンライン平方分割のようなデータ構造が確か存在します
  - Two-Level Tree という巡回セールスマン問題の中程度のインスタンス用のデータ構造です
  - Fredman, M. L., Johnson, D. S., McGeoch, L. A., & Ostheimer, G. (1995). Data structures for traveling salesmen. Journal of Algorithms, 18(3), 432-479.
- もしかしたら今回の手法も似た形でオンライン化が可能かもしれません (が、速度が実用に耐えない予感がします)
東京工業大学プログラミングコンテスト2022 - AtCoder
- https://atcoder.jp/contests/ttpc2022
- 2015と2019もよろしくお願いします