全方位木DPに入門してABC220のF問題を解く

## はじめに

全方位木DPは、ノード数$N$の木のすべてのノード$i$について$i$を根としたときに知りたい値を$O(N)$で求められるアルゴリズムです。何も知らないと$i$でループし、$i$を根とした場合の探索をして$O(N^2)$なので、全方位木DPによってとても速くなります。

[ABC220のF問題](https://atcoder.jp/contests/abc220/tasks/abc220_f)を例にして、全方位木DPの気持ちを理解し、実際にACすることを確認します。

この記事では、以下の木を例にします。この木は、以下の入力に対応します。

```
6
1 2
2 3
2 4
4 5
4 6
```

また, 説明の際に以下の記号を用います。

* $p_i$: ノード$1$を根としたときの、ノード$i$の親ノード。例えばこの木では$p_4=2$
* $C_i$: ノード$1$を根としたときの、ノード$i$の子ノード集合。例えばこの木では$C_4=\\{5, 6\\}, C_5=\emptyset$
* $A_i$: ノード$i$に隣接するノードの集合。例えばこの木では$A_4=\\{2, 5, 6\\}$
* $T_i$: ノード$i$を根としたときの, ノード$i$の部分木に含まれるノード集合。例えばこの木では$T_2=\\{2, 3, 4, 5, 6\\}$

![木.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/656409/9a73ca4d-09fa-b0a9-7d5f-075fc6d5f6a4.png)

## 問題概要

$N$頂点の重みなし木が与えられます。頂点$i, j$間の距離を$dis(i, j)$とするとき, $i=1,\ldots,N$のそれぞれについて$\displaystyle \sum\_{j=1}^ndis(i, j)$を出力せよ。



## 木DPについて
木DPは根を$r$に固定した木で、なんらかの配列$\text{dp}$の$i$番目の要素$\text{dp}_i$にノード$i$を根とする部分木の計算結果を保存していき、最終的に木全体（$\text{dp}_r$の値）での値を求めます。

ここでは$r=1$として考えます。木DPを用いて、$\sum_{j=1}^ndis(1, j)$を求めます。（$i=1$の場合に相当します。）

ここでは以下の2種類の値を求めていきます。（この理由は後ほど説明します。）

* $\\displaystyle\text{count}_i =\\left|T_i\right|$: $i$を根とする部分木のノード数

* $\\displaystyle\text{dis_sum}_i=\sum\_{j\in T_i}dis(i, j)$: $i$を根とする部分木の各ノード$j$について$dis(i, j)$を合計した値


$i=1$のとき、部分木が木全体になることから, $\text{dis_sum}_1$が求める値になることがわかります。

1. $i=3, 5, 6$のとき,
    部分木はただ1つのノードなので、$\text{count}_i=1,\text{dis_sum}_i=0$です。

    |$i$|$1$|$2$|$3$|$4$|$5$|$6$|
    |:---|---|---|---|---|---|---|
    |$\text{count}_i$|||$1$||$1$|$1$|
    |$\text{dis_sum}_i$|||$0$||$0$|$0$|

  ![dp_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/656409/351cee66-f65a-35f1-1ad1-cc3cafc18064.png)

2. $i=4$のとき,
    部分木のノード数は, 子ノードの部分木のノード数の総和に1（自分自身）を加えた数になります。
    $\text{count}_4=\text{count}_5+\text{count}_6+1=3$
    また, 4以外のノードへは5または6を経由するので、4-5, 4-6の辺をノード数だけ通過し、それだけ$dis(i, j)$の総和も増加します。これが、$\text{count}_i$と$\text{dis_sum}_i$の2種類を計算する理由です。
    ```math
    \begin{align*}
    \text{dis_sum}_4&=dis(4, 4) + \sum_{j\in T_5}dis(4, j) + \sum_{j\in T_6}dis(4, j)&\\
    &=0+\sum_{j\in T_5}(1+dis(5, j)) + \sum_{j\in T_6}(1+dis(6, j))&\\
    &=|T_5|+|T_6|+\sum_{j\in T_5}dis(5, j)+\sum_{j\in T_6}dis(6, j)&\\
    &=\text{count}_5+\text{count}_6+\text{dis_sum}_5+\text{dis_sum}_6
    \end{align*}
    ```

    |$i$|$1$|$2$|$3$|$4$|$5$|$6$|
    |:---|---|---|---|---|---|---|
    |$\text{count}_i$|||$1$|$3$|$1$|$1$|
    |$\text{dis_sum}_i$|||$0$|$2$|$0$|$0$|

![dp_4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/656409/6296f3df-f609-b349-ceba-875adf215a7c.png)


2. $i=2, 1$のとき, 同様にして以下のように表が埋まります。 

    |$i$|$1$|$2$|$3$|$4$|$5$|$6$|
    |:---|---|---|---|---|---|---|
    |$\text{count}_i$|$6$|$5$|$1$|$3$|$1$|$1$|
    |$\text{dis_sum}_i$|$11$|$6$|$0$|$2$|$0$|$0$|

## 根を変更する

$i=1$のときの答えが出せました。ここから, $i=2,\ldots,n$のときの答えを求めます。

$i=2,\ldots,n$のときにはノード$i$には親$p_i$がいるので、これを利用します。

処理のイメージとしては以下の流れです。

1. ノード$i$とノード$p_i$の間の辺を削除して、別々の木として計算する
1. ノード$p_i$を$i$の子ノードに加えて再計算する
1. 上記の1,2を親から順に行う

$i=2$のときの処理は以下のようになります。

![rerooting.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/656409/a371bcd6-5a33-2e08-5fd6-f65072cd06dd.png)

もともと2の子ノードだった3, 4は最初に求めていたcount, dis_sumが使えますが、親ノードだった1については再計算が必要です。

ノード1を子ノードとして扱うには、1-2を結ぶ辺を削除しているので、
* $\displaystyle \text{count}'_1 = 1+\sum\_{j\in A_1, j\ne 2}\text{count}_j$
* $\displaystyle \text{dis\_sum}'_1 = \sum\_{j\in A_1, j\ne 2}(\text{count}_j+\text{dis\_sum}_j)$
が必要です。

ここでループして計算してしまうと全体で最悪$O(N^2)$になってしまうのですが、累積和のような方法を使うことで$O(N)$に抑えることができます。これで$i=2$での値がわかりました。

$i=2$の結果から同様にして$i=3, 4$の場合がわかり,$i=4$の結果から同様にして$i=5, 6$の結果がわかります。

## 全方位木DPまとめ（ABC220 F問題）
1. 辺の情報を受け取って$A_i$, $p_i$, $C_i$を構築する。
    * $p_i$は深さ優先探索などで求まる。
    * $C_i$は$A_i - \\{p_i\\}$により求まる。
1. 木DPで$i=1$のときの値を求める。
    1.葉から順にループする。ノード$j$について、$\text{count}_j$, $\text{dis\_sum}_j$を求める。
    * $\displaystyle \text{count}_j=1+\sum\_{c\in C_j}\text{count}_c$
    * $\displaystyle \text{dis\_sum}_j=\sum\_{c\in C_j}(\text{count}_c+\text{dis\_sum}_c)$
1. 根から順にループして、 $i=2,\ldots,N$のときの値を求める。
    1. $j\in A_i$のそれぞれについて、$\text{count}'_j$, $\text{dis\_sum}'_j$を求める。
        * $j\in C_i$のとき, $\text{count}'_j = \text{count}_j$, $\text{dis\_sum}'_j=\text{dis\_sum}_j$
        * $j=p_i$のとき, $\displaystyle \text{count}'_j=1+\text{count\_without}_j(i)$, $\text{dis\_sum}'_j=\text{dis\_sum\_without}_j(i)$
    1. $\displaystyle \sum\_{j\in A_i}(\text{count}'_j+\text{dis\_sum}'_j)$が、求めたい値である。
    1. $j\in A_i$のそれぞれについて、$\displaystyle \text{count\_without}_i(j)=\sum\_{c\in A_i, c\ne j}\text{count}'_c$, $\displaystyle \text{dis\_sum\_without}_i(j)=\sum\_{c\in A_i, c\ne j}\text{dis\_sum}'_c$を求める。

        
## 全方位木DPまとめ（一般）

1. 辺の情報を受け取って$A_i$, $p_i$, $C_i$を構築する。
    * $p_i$は深さ優先探索などで求まる
    * $C_i$は$A_i - \\{p_i\\}$により求まる
1. 木DPで$i=1$のときの値を求める。
    1.葉から順にループする。ノード$j$について、$\text{dp}_j$を求める。
    * $\displaystyle \text{dp}_j=f\_{\text{after}}\left(\underset{c\in C_j}{F\_{\text{merge}}}\bigl[f\_{\text{before}}(\text{dp}_c, j, c)\bigr], j\right)$
    * $f\_{\text{before}}(\text{dp}_c,\text{parent}, \text{child})$:子ノードのdpの値を集約前に加工する関数。辺の情報を付加するときなど。
    *  $F\_{\text{merge}}(v_1, v_2,\ldots)$:子ノードの情報を集約する関数。合計や最大など、結合側が成立する演算である必要がある。（モノイド）
        * 単位元を$I$とする。  
    *  $f\_{\text{after}}(\text{merged},\text{node})$:集計後に加工する関数。ABC220Fでは$\text{count}_j$の総和をとったあとに1を加算したり, $\text{dis\_sum}_j$に$\text{count}_j$を加算したりしている。
1. 根から順にループして、 $i=2,\ldots,N$のときの値を求める。
    1. $j\in A_i$のそれぞれについて、$\text{dp}'_j$を求める。
        * $j\in C_i$のとき, $\text{dp}'_j = \text{dp}_j$
        * $j=p_i$のとき, $\displaystyle \text{dp}'_j=\text{dp\_without}_j(i)$
    1. $\displaystyle f\_{\text{after}}\left(\underset{c\in A_i}{F\_{\text{merge}}}\bigl[f\_{\text{before}}(\text{dp}'_c, i, c)\bigr], i\right)$が、求めたい値である。
    1. $j\in A_i$のそれぞれについて、$\displaystyle \text{dp\_without}_i(j)=f\_{\text{after}}\left(\underset{c\in A_i, c\ne j}{F\_{\text{merge}}}\bigl[f\_{\text{before}}(\text{dp}'_c, i, c)\bigr], i\right)$を求める。
        * $c\in A_i, c\ne j$は $c\in A_i, c < j$または $c\in A_i, j < c$なので,$j$未満で集約した値と, $j+1$以上で集約した値を集約します。
        * 厳密には次の手順を行います。$A_i$を配列表現する。配列$\text{mf}, \text{ml}$を以下のようにして求める。（それぞれmerged_first, merged_lastの略です）

        ```math
        \begin{align*}
        \text{mf}[0] &= I,&\\
        \text{mf}[k] &= F_{\text{merge}}\bigl( \text{mf}[k-1], f_{\text{before}}(\text{dp}'_{A_i[k]}, i, A_i[k])\bigr)\ \ (k=1,\ldots, |A_i|)&\\
        \text{ml}[|A_i|+1] &= I,&\\
        \text{ml}[k] &= F_{\text{merge}}\bigl( \text{ml}[k+1], f_{\text{before}}(\text{dp}'_{A_i[k]}, i, A_i[k])\bigr)\ \ (k=|A_i|,\ldots, 1)&\\
        \end{align*}
        ```

        とすると, 

        ```math
        \begin{align*}
        \text{mf}[k] &= \displaystyle \underset{0\leq m\leq k}{F_{\text{merge}}}\bigl[f_{\text{before}}(\text{dp}'_{A_i[m]}, i, A_i[m])\bigr]&\\
        \text{ml}[k] &= \displaystyle \underset{k\leq m\leq |A_i|}{F_{\text{merge}}}\bigl[f_{\text{before}}(\text{dp}'_{A_i[m]}, i, A_i[m])\bigr]
        \end{align*}
        ```

        が保存される。$\displaystyle \underset{c\in A_i, c\ne A_i[k]}{F\_{\text{merge}}}\bigl[f\_{\text{before}}(\text{dp}'_c, i, c)\bigr]=F\_{\text{merge}}\bigl(\text{mf}[k-1], \text{ml}[k+1]\bigr)$より、 $\text{dp\_without}_i(j)$を求めます。