ARC230 A - Meeting on Tree

考え方

まず、集合場所をどこにすべきかを考える。

集合地点を仮にある頂点 $P$ とする。
各リスが木の辺を伝ってくる経路は $1$ つである。
このとき、集合地点に接続されたある辺 $Q\to P$ を参加リスの過半数が通るとしよう。
この場合、集合地点を $P$ から $Q$ に変更した方がコストが小さくなる。

つまり、最善の集合地点とは、参加リスの過半数が通るような辺がなくなる地点である。
各辺、少ない側から多い側への一方通行としてそれに従ってみんな移動すればいいというわけである。
（ちょうど半数ずつだった場合は、どちらに決めてもコストは同じ）

ということを考えると、最善の場合のコストは、辺ごとに計算するのがよさそう。
具体的には、選び方 $2^N-1$ 通りそれぞれ、辺から見て参加リス数の少ない方の合計を求めればよい。

ということで、ある辺に注目し、少ない側に $L$ 匹いるとして、これを左側と呼ぶことにする。
多い側には $N-L$ 匹のリスがいて、これを右側と呼ぶことにする。
ちょうど半分なら、好きな方を左側にしてよい。
コード的には、この $L$ の値は単純な木DPで求まる。

参加リスの選び方 $2^N-1$ 通り、左側からの参加リス数を合計すると、延べ $2^{N-1}L$ 匹。
これは、各リスの参加パターンが $2^{N-1}$ 通りであることからわかる。
ここから、実は右側の方が参加リス数が少ないパターンについて、差分を引けばよい。

右側の方が参加リスが $k$ 匹少ないような選び方の総数は、以下のように求められる。
$$
\sum^{L-k}_{i=0} {}_L\mathrm{C}_{i+k} {}_{N-L}\mathrm{C}_{i}
= \sum^{L-k}_{i=0} {}_L\mathrm{C}_{L-i-k} {}_{N-L}\mathrm{C}_{i}
= {}_N\mathrm{C}_{L-k}
$$

最後の変形は、ヴァンデルモンドの恒等式を用いた。
$L$ 個から $L-i-k$ 個、$N-L$ 個から $i$ 個選ぶ操作の総和は、ただ $N$ 個から $L-k$ 個選ぶだけである。

よって、これに匹数の差をかけて和を取ることで差分合計が求められる。
$$
\begin{aligned}
\sum^{L}_{k=0} k {}_N\mathrm{C}_{L-k}
= \sum^{L}_{k=0} (L-k) {}_N\mathrm{C}_{k}
= L \sum^{L}_{k=0} {}_N\mathrm{C}_{k} - \sum^{L}_{k=0} k {}_N\mathrm{C}_{k}
\end{aligned}
$$

後のコードの書きやすさのために、$L-k$ を $k$ に取り直した。
これら $2$ つの和は、$N$ の値は問題を通して一定であり、和の中身に $L$ が含まれない。
よって、階乗、二項係数、累積和、$L$ の値の $4$ つの前処理のもと、この値は辺 $1$ つあたり $O(1)$ で求まる。
この差分を前述の $2^{N-1}L$ から引けば、その辺を通る分のコスト合計である。

$4$ つの事前処理は全て $O(N)$ ででき、その後 $O(1)$ 処理が $N-1$ 回なので、全体計算量は $O(N)$ である。

入力例2での動作

入力を受け取る。

n: 5
edges:
(1, 2)
(2, 3)
(3, 4)
(3, 5)

入力例2の木を頂点 $1$ を根として根付き木にする。
頂点 $1,2,3$ が順につながり、頂点 $3$ の子が頂点 $4,5$ となる。

まず、$2^{N-1}=2^4=16$ である。
また、${}_5\mathrm{C}_0=1$, ${}_5\mathrm{C}_1=5$, ${}_5\mathrm{C}_2=10$ なので、必要な累積和は以下のようになる。

葉から順に部分木の頂点数を求める。
頂点 $4,5$ の部分木サイズはそれぞれ $1$、頂点 $3$ は $3$、頂点 $2$ は $4$、頂点 $1$ は $5$ となる。

各頂点と親を結ぶ辺について、辺を切ったときの少ない側のリスの総数 $L$ は以下のようになる。

$L=1$ の辺のコストへの寄与は $1\times 16-1\times 6+5=15$ となる。

$L=2$ の辺のコストへの寄与は $2\times 16-2\times 16+25=25$ となる。

したがって、$L=1$ の辺が $3$ 本、$L=2$ の辺が $1$ 本なので、答えは $15\times3+25=70$ となる。

注意点

二項係数は $998244353$ を法として計算する。
加算・乗算のたびに剰余を取り、逆階乗を求めるための除算は逆元を用いる。

別解

特になし。