−目次
AtCoder Beginner Contest 220 E,F問題メモ
E - Distance on Large Perfect Binary Tree
問題
- 2N−1 個の頂点からなる完全二分木
- 頂点 1 を根とし、頂点 i の子は 2i,2i+1
- 頂点 (i,j) の組で、その最短距離が D となるものの個数を mod998244353 で求めよ
- 1≤N≤106
解法
パターンが同じになるものを見つけて、まとめて計算する。
D=1 の場合、答えは辺数×2 =2N+1−4 となる。
以下、D≥2 とする。
完全二分木なので、深さが同じ頂点以下の部分木は形状が全く一緒となる。
当然、その中だけで取れる答えの対象ペアの個数も一緒となる。
一方、深さが異なる頂点の部分木同士はまとめて考えることはできないため、深さでグループ分けするのがよさそう。
深さ d に位置する頂点(v とする)毎に、
「v が2頂点を結ぶパスに含まれ、かつそれが最も浅い頂点である」組の個数を数えると、重複せずに数え上げられる。
以下の2つに分けて考える。
- A.vを一端、もう一端をそこより D 深い頂点とする組
- B.v の左の子孫から1点、右の子孫から1点選び、距離が D となる2頂点の組
N=5 d ① 0 ,-----' `-----, ② ③ 1 / \ / \ ④ ⑤ ⑥ ⑦ 2 /\ /\ /\ /\ ⑧ ⑨ ⑩ ⑪ ⑫ ⑬ ⑭ ⑮ 3 /\ /\ /\ /\ /\ /\ /\ /\ ⑯⑰⑱⑲⑳㉑㉒㉓㉔㉕㉖㉗㉘㉙㉚㉛ 4
A. まっすぐ降りる
上記の例で、例えば D=3 のとき、②にとっての (②,⑯), (②,㉒) などが該当する。
d+D≥N なら、D だけ深い頂点は存在しない。
d+D<N の場合、1頂点につき 2D のペア対象が存在する。(②に対する⑯~㉓)
今回の問題では (i,j),(j,i) は別に数えるので2倍、さらに深さ d には 2d 個の頂点があるので、
全部で 2d+D+1 となる。
B. 左右から1点ずつ
上記の例で、例えば D=3 のとき、②にとっての (⑧,⑤), (④,⑪) などが該当する。
左の頂点を深さ d+l の頂点とすると、距離 D にするためには右は深さ d+(D−l) となる。
②にとって、自身より l 深い左の頂点は 2l−1 個、D−l 深い右の頂点は 2D−l−1 個で、 この組合せは 2D−2 個となる。
つまり、l に関わらず、どの深さとどの深さをペアにしても組合せは常に 2D−2 で固定である。いい感じ。
では何通りの l が取れるのかだが、
- 両方を最大の深さにしても、D に足りない場合
- 1つも取れないのでスキップ
- 最小の 1 から最大の D−1 までフルで取れる場合
- D−1 通り
- l が小さすぎると相方が深すぎて存在しなくなる場合
- 相方の深さがMAX d+D−l=N−1 のとき、l が取れる最小値 d+D−N+1
- 従って d+D−N+1~N−1 の 2N−d−D−1 通り
この2つの小さい方が、l が取れるパターン数(p とする)となる。
A.と同様、反転させたペアで2倍、深さ d には 2d 個の頂点があるので、
全部で p×2d+D−1 となる。
d=0~N−2 まで足し合わせた結果が答え。
F - Distance Sums 2
問題
- N 頂点の木、i 番目の辺は ui,vi を結ぶ無向辺
- 各頂点について、他の全頂点との距離の総和を求めよ
- つまり、i=1,2,...,N について、N∑j=1dist(i,j) を求めよ
- 2≤N≤2×105
解法
主客転倒で考える。
1本の辺が、どの頂点にどれだけ寄与するのか考える。
... ○--, e ,--○ ○-○---○-○-○ ... ○----' `--○ └────┘ └──────┘ A B
こうなったとき、辺 e は、
- A 側の頂点にとって、B 側のどれか1個とペアになったときに1回ずつ、計 |B| 回答えに加算される
- B 側の頂点にとって、A 側のどれか1個とペアになったときに1回ずつ、計 |A| 回答えに加算される
つまり、A 側の頂点に全て |B| を加算し、B 側の頂点に全て |A| を加算すればよい。
これを全ての辺に対して行えば、答えとなる。
加算には、木の累積和を用いる。
適当に頂点1を根としておいて、まず1回目のDFSで各頂点の部分木のサイズ size[v] と深さを求める。
次いで、以下の配列を用意する。最初は全て0。
- Acc[v]=v 以下の部分木の全頂点に、Acc[v] の値を加算する
各辺について、両端の頂点の内、深さが浅い方(親)を u、深い方(子)を v とすると、
- v 側の全頂点に、u 側のサイズ N−size[v] を加算する
- Acc[v]+=N−size[v]
- u 側の全頂点に、v 側のサイズ size[v] を加算する
- これは、まず全頂点に size[v] を加算し、v 以下の頂点ではそれを打ち消せばよい
- Acc[1]+=size[v]
- Acc[v]−=size[v]
最後に2回目のDFSで、Acc を根から順に子に伝播させていけばよい。