AtCoder Beginner Contest 160 F - Distributing Integers

F - Distributing Integers

解法

「頂点に整数を書く」だと少し図示しづらいので、代わりに「頂点番号の書いたボールを処理する順番に並べる」として考えることにします。

最初にボールを置く頂点を根付き木の根とします。まずは根を頂点 $1$ に固定して問題を解くことにしましょう。以下のような木DPを組みます。

$dp\lbrack i \rbrack =$ グラフが頂点 $i$ 以下の部分木だけしかないと仮定した時に、最初に頂点 $i$ のボールを置き、ルールに沿って部分木全ての頂点のボールを置くときの並べ方の数

これを葉から順番に計算していきます。根 $1$ についての答えは $dp\lbrack 1 \rbrack$ になります。

$dp\lbrack i \rbrack$ の計算方法を考えます。頂点 $i$ の子を $c_{1}, ..., c_{K}$ と書くことにします。この時のボールの並べ方は以下のように考えることができます。

f:id:betrue12:20200329001308p:plain

つまりそれぞれの子以下のボールの並べ方 $dp\lbrack c_{k}\rbrack$ を全て掛け合わせ、さらにそれらを混ぜる方法の数を掛けることで $dp\lbrack i \rbrack$ を求めることができます。

頂点 $v$ 以下の部分木に含まれる頂点の個数を $n_{v}$ と書くことにします。具体的な計算式は、階乗を用いて

$\displaystyle dp\lbrack i \rbrack = (\prod_{c_{k}}dp\lbrack c_{k}\rbrack) \times \frac{(n_{c_{1}} + \cdots + n_{c_{K}})!}{n_{c_{1}}! \times \cdots \times n_{c_{K}}!}$

とまとめて計算することもできますが、後々使うので「子を1つずつマージする」ような計算方法も確認しておきます。

初期値として、 $dp\lbrack i\rbrack = 1, n_{i} = 0$ とします。 $n_{i}$ はマージ済みの子孫の頂点数として扱い、頂点 $i$ 自身は最後になるまで加算しません。ここに子 $c_{k}$ をマージするときには、

$n_{i}$ に $n_{c_{k}}$ を加算する。
$dp\lbrack i\rbrack$ に $dp\lbrack c_{k}\rbrack \times {}_{n_{i}}C_{n_{c_{k}}}$ を乗算する。

という処理をすれば良いです。これを全ての子について行い、最後に頂点 $i$ 自身のぶんとして $n_{i}$ に $1$ 加算すれば完了です。この計算結果は先ほどの階乗で表現した結果と同じになります。

全頂点を根とする場合に拡張

さて、これで1つの頂点を根とする場合は解くことができました。実際には全ての頂点を根とした場合の答えを求めないといけないので、俗に「全方位木DP」と呼ばれるテクニックを用います。

全方位木DPは、多くの場合2回のDFSで実装できます。1回目は先ほどのように、根を1つに固定した時の計算を葉から順番に行います。2回目のDFSでは根から順番に、全ての頂点について「その頂点が根であるときの答えの値」を求めることになります。このときに「根をずらす」という処理を行います。

今回の問題で具体的に説明します。まず図でイメージを表現するとこんな感じです。

f:id:betrue12:20200329001550p:plain

$c_{1}$ を根としたときの答えを求める方法を考えます。ここで、元の木における $c_{1}$ 以下の部分木を並べる方法は既に $dp\lbrack c_{1}\rbrack$ として求めています。そのため、そこに $c_{1}$ の親だった頂点 $i$ を $c_{1}$ の子としてマージするような係数を掛ければ良いです。

そのため「 $c_{1}$ の子としての」 $dp\lbrack i \rbrack$ に相当する値が必要になります。1回目のDFSでは $c_{1}$ が子であるときの $dp\lbrack i \rbrack$ を求めているので、今度は $dp\lbrack i \rbrack$ に対して子だった $c_{1}$ のマージ操作を逆戻しにするような係数を掛ければ良いです。

このマージ操作やマージの逆操作は、先ほど計算した「子を1つずつマージする」ときの係数あるいはその逆数を掛け算することに相当します。計算がごちゃごちゃしますがしっかり図に描いて整理しましょう。頂点 $i$ や $c_{1}$ を根として扱っているときには、（自身を除く）子孫の頂点数は $N-1$ であることを意識すると少し整理しやすいです。具体的な係数は以下のようになります。

$dp\lbrack i \rbrack$ から $c_{1}$ のマージ操作を逆戻しにするときには、部分木の $n_{c_{1}}$ 個の頂点をマージする操作を戻すので、係数は $dp\lbrack j \rbrack \times {}_{N-1}C_{n_{c_{1}}}$ の逆数となる。
$dp\lbrack c_{1}\rbrack$ に子としての頂点 $i$ をマージするときには、マージする部分木の頂点数は $N-n_{c_{1}}$ となるので、「 $c_{1}$ の子としての」 $dp\lbrack i \rbrack$ の値を $u$ として係数は $u \times {}_{N-1}C_{N-n_{c_{1}}}$ となる。