geam1113’s diary

主にAtcoderのコンテストの備忘録を書きます。AtCoder緑と水色を行ったり来たりしています。

ABC220F Distance Sums 2

ABC 全方位木DP

与えられた木を、頂点1を根とする根付き木とみなします。

辺 $e=(u,v)$ で結ばれた頂点 $u,v$ を考えます。頂点 $u$ の方が根に近いものとします。

根からDFSをすると、 $v$ での探索が終了した後に $u$ に戻ります。

$v$ を根とする部分木について、 $v$ の子孫 $w_1,...w_k$ についての $dis(v,w_1)+...+dis(v,w_k)$ が求まっているとします。

$u$ から見た時、 $w_1,...w_k$ までの距離は、 $v$ までの距離に $dis(u,v)=1$ が追加されたと考えることができます。式で表すと、 $dis(v,w)=dis(u,v)+dis(v,w)$ です。

$dis(u,v)$ が何回足されるかを考えると、 $v$ の子孫の個数+1( $v$ の分)となります。これは $v$ を根とする部分木の頂点数に一致します。

したがって、根である1からの各頂点へ距離の総和は、以下の木DPにより求まります。

$dp1[u]$ :=頂点1を根とする根付き木において、頂点 $u$ を根とする部分木の頂点数

$dp2[u]:=$ 頂点1を根とする根付き木において、頂点 $u$ を根とする部分木の、 $u$ とその子孫 $w$ との距離 $dis(u,w)$ の総和

遷移を考えます。

・ $dp1$

$u$ の子 $v_1,...,v_k$ について、 $dp1[v_1],...,dp1[v_k]$ が求まっているとすると、

$u$ を根とする部分木の頂点数は、子を根とする部分木の頂点数+ $u$ 自身となるので、以下の通りです。

$dp1[u] = dp1[v_1]+...+dp1[v_k]+1$

・ $dp2$

$u$ の子 $v_1,...,v_k$ について、 $dp2[v_1],...,dp2[v_k]$ が求まっているとすると、 $u$ からすべての子孫への距離は、子である $v_i$ からの距離+ $dis(u,v)(=1)$ となり、 $dis(u,v)$ は、 $v_i$ を根とする部分木の頂点数分足されるので以下の通りです。

$dp2[u] = dp2[v_1]+...+dp2[v_k]+dp1[v_1]+...+dp1[v_k]$

このDPをDFSしながら解くと、根である頂点1からの各頂点の距離の総和は求まります。

しかし、今回「すべてを根とした時」について総和を求める必要があります。

この時に使えるのが全方位木DPです。

全方位木DPの詳細な解説は検索するとわかりやすい記事が出てくるので割愛します。

簡単に述べると、

(1) 頂点1を根とする木DPをする(DFSする)。

(2) 根でない頂点 $v$ については、1回目のDFSで親である頂点 $u$ を根とする $v$ でない側の部分木についてのみ、値が計算されていない

(3) そこで、 $u$ のうち、 $v$ が接続しているところを除いた部分木のDPの値を計算する

(4) 計算したDPの値を $v$ に渡す

(3),(4)のアルゴリズムとしては、 $u$ の子 $v_1,...,v_k$ について、 $dp_p=dp[v_1]+...+dp[v_{i-1}]+dp[v_{i+1}]+...+dp[v_k]$ を計算して、 $v_i$ に $dp_p$ を渡します。この処理を頂点1から順に行なっていきます。(これはBFSになるらしいです。)

といったことをします。

全方位木DPでは、頂点 $u$ から見たときの頂点 $v$ を根とした時の部分木、というふうに見るため、以下のようなDPとなります。

$dp1[u][v]:=$ 頂点 $u$ を親としたときの、頂点 $v$ を根とする部分木の頂点数

$dp2[u][v]:=$ 頂点 $u$ を親としたときの、頂点 $v$ を根とする部分木の、 $v$ とその子孫 $w$ との距離 $dis(v,w)$ の総和

$u$ があるかないかで、遷移はほぼ同じです。

全方位木DPでは、辺 $e=(u,v)$ に対し、始点側を $u$ とすると、

・ $v$ のDP情報は $u$ が保存している

という点で、通常の木DPと異なると思います。

また、

(1) $v$ が接続する頂点のDPの値を統合しつつ( $u$ を除く)、何らかの計算を行い、

(2) その値から、 $v$ を含めたDPの値を計算し直し、

その計算値を $u$ 側に渡すことが必要になります。

(1)の処理だけでは、自分を含めた計算値になっていない場合があることにも留意します。

上記における、(1)の子の情報を統合して計算するのをop、(2)の自分を含めたDPの値を計算し直すのをadd_rootとして定義します。

今回の問題で実装を考えます。

$u$ と接続する $v_1,...,v_k$ について考えます。

opでは $dp1,dp2$ 共に、単純に子の情報を足し合わせていきます。

すると、最終的に

$dp1_u=dp1[u][v_1]+...+dp1[u][v_k]$

$dp2_u=dp2[u][v_1]+...+dp2[u][v_k]$

が計算された状態になっています。

この状態から、 $u$ に関する情報を計算し直すadd_rootでは、

$dp2_u$ に $dp1_u+1$ を加え、 $dp1_u$ に1を加えると $u$ についての正しいDPの値となります。

もう1点気をつけたいのが、全方位木DPでは、自分に接続する $i$ 番目の頂点の情報を除くために、

$dpLeft_i = op(dp_1,...,dp_{i-1})$
$dpRight_i = op(dp_i,...,dp_k)$

を予め計算値しておき、

$op(dpLeft_i, dpRight_{i+1})$ によって、目的の値を得ます。

opはDFSでの子の情報の統合と、上記計算の3箇所で使用され、いずれも正しく値を計算できなければならないです。(コンテスト中、ここの実装でつまずいていました。)

自分の間違いとして、 $op$ について、

$op(dp1_i,dp1_j)=dp1_i+dp1_j$

$op(dp2_i,dp2_j)=dp2_i+dp2_j+dp1_j$

としていましたが、これだと、 $op(dpLeft_i, dpRight_{i+1})$ を計算した時、不具合がおきます。

提出コード:https://atcoder.jp/contests/abc220/submissions/26256944

追記

ABC222Fの解説が全方位木DPについてよく書かれています。

https://atcoder.jp/contests/abc222/editorial/2749