全方位木DP(ReRooting)の抽象化について

「 $N$ 頂点の木において、 $v=1, 2, \dots, N$ を根としたときの○○をそれぞれ求めよ」というタイプの問題では、全方位木DP(ReRooting)が有効に働く場合が多いです。
根を1つ固定したときにこの問題が高速に解けるならば、全方位木DPを用いることで同じ計算量で全ての根についての問題を解くことが出来ます。
この記事では主に実装方法について述べるので、アルゴリズム自体の詳しい解説は他の方のブログを参照してください。

また、この記事を通して、根を固定して考えている際に「 $v$ を根とする部分木」と言った場合、親方向に辺が伸びるようなものは考えません。

全方位木DPの抽象化

この記事の本題です。
適当な頂点を根とした場合の木DPを考えます。
$T$ を何らかの集合とし、 $dp_v \in T$ を $v$ を根とする部分木の何らかの場合の数/判定等とします。
このとき、木DPの遷移が

2変数関数 $f, g: T \times \mathbb{N} \rightarrow T$
可換モノイド $merge: T \times T \rightarrow T$
$v$ の子の集合 $ch(v)$

を用いて

$\displaystyle dp_v = g(merge(f(dp_{c_1},c_1), f(dp_{c_2},c_2), \dots, f(dp_{c_k},c_k)),v) \ (c_i \in ch(v))$

と表せるならば、後述する全方位木DPのテンプレートが使えます。モノイドが非可換な場合でもあらかじめマージの順番が決まっている場合は、隣接リストを適当にソートすることで対応できます。
$f,g$ は第2引数に頂点番号を取りますが不要なことも多いです。 $merge$ は2回目のDFSで左右から累積を取る際に結合則が必要であるため、モノイドでなければなりません。
これだけではわかりにくいと思うので、後程いくつか例を挙げます。

ソースコード

template <typename T>
struct ReRooting {
  using F = function<T(T, int)>;
  using F2 = function<T(T, T)>;
  int V;
  vector<vector<int>> G;
  vector<vector<T>> dp;
  // dp_v = g(merge(f(dp_c1,c1), f(dp_c2,c2), ..., f(dp_ck,ck)), v)
  F f, g;
  F2 merge;
  T mi;  // identity of merge

  ReRooting() {}
  ReRooting(int V, F f, F2 merge, T mi, F g)
      : V(V), f(f), merge(merge), mi(mi), g(g) {
    G.resize(V);
    dp.resize(V);
  }

  void read(int idx = 1) {
    int a, b;
    for (int i = 0; i < V - 1; ++i) {
      cin >> a >> b;
      a -= idx, b -= idx;
      G[a].emplace_back(b);
      G[b].emplace_back(a);
    }
  }

  void add_edge(int a, int b) {
    G[a].emplace_back(b);
    G[b].emplace_back(a);
  }

  T dfs1(int p, int v) {
    T res = mi;
    for (int i = 0; i < G[v].size(); ++i) {
      if (G[v][i] == p) continue;
      dp[v][i] = dfs1(v, G[v][i]);
      res = merge(res, f(dp[v][i], G[v][i]));
    }
    return g(res, v);
  }

  void dfs2(int p, int v, T from_par) {
    for (int i = 0; i < G[v].size(); ++i) {
      if (G[v][i] == p) {
        dp[v][i] = from_par;
        break;
      }
    }
    vector<T> pR(G[v].size() + 1);
    pR[G[v].size()] = mi;
    for (int i = G[v].size(); i > 0; --i)
      pR[i - 1] = merge(pR[i], f(dp[v][i - 1], G[v][i - 1]));
    T pL = mi;
    for (int i = 0; i < G[v].size(); ++i) {
      if (G[v][i] != p) {
        T val = merge(pL, pR[i + 1]);
        dfs2(v, G[v][i], g(val, v));
      }
      pL = merge(pL, f(dp[v][i], G[v][i]));
    }
  }

  void calc(int root = 0) {
    for (int i = 0; i < V; ++i) dp[i].resize(G[i].size());
    dfs1(-1, root);
    dfs2(-1, root, mi);
  }

  T solve(int v) {
    T ans = mi;
    for (int i = 0; i < G[v].size(); ++i)
      ans = merge(ans, f(dp[v][i], G[v][i]));
    return g(ans, v);
  }
};

変数f, g, mergeが先程述べた関数に対応しています。miはmergeの単位元です。
関数calc()を呼んだ後関数solve(i)を呼ぶことで、iを根とする場合の解が求まります。

追記：コンストラクタでE, dpをresizeするよう変更しました。（Codeforcesのようにテストケース式の場合、毎回MAX_V個のvectorを生成するとTLEするため）
さらに追記(2020/09/27)：便利関数を追加し、その他修正を行いました。

使用例

例1. EDPC-V Subtree

問題リンク

V - Subtree

問題概要

$N$ 頂点の木において、 $v=1, 2, \dots, N$ を根とする部分木の個数（を $10^9+7$ で割った余り）をそれぞれ求めよ。

解法

適当な頂点を根として考えます。（以下その他の問題でも同様）
$dp_v$ を $v$ を根とする部分木の個数とします。
DPの遷移は

$\displaystyle dp_v = \prod_{c \in ch(v)}(dp_c + 1)$

となります。
よって $f(a,v) = a+1, \ merge(a,b) = ab, \ g(a,v) = a$ とすればよいです。

提出コード：
atcoder.jp

例2. Codeforces Round #627 (Div. 3) F

問題リンク

Problem - F - Codeforces

問題概要

各頂点が白または黒のいずれかで塗られた $N$ 頂点の木がある。
$v=1, 2, \dots, N$ を含む部分木の（白の頂点数）-（黒の頂点数）の最大値をそれぞれ求めよ。

解法

$dp_v$ を $v$ を根とする部分木（ $v$ を必ず含む）の（白の頂点数）-（黒の頂点数）の最大値とします。
$c \in ch(v)$ について、 $dp_c$ が正ならその部分木を $v$ に接続し、そうでなければ接続しないのが最適です。
したがって、 $w_i$ を $i$ 番目の頂点が白のとき1、黒のとき-1とすると、DPの遷移は

$\displaystyle dp_v = \sum_{c \in ch(v)}(max(dp_c,0)) + w_v$

となります。
よって $f(a,v) = max(a,0), \ merge(a,b) = a+b, \ g(a,v) = a+w_v$ とすればよいです。

提出コード：
codeforces.com

例3. ABC160-F Distributing Integers

問題リンク

F - Distributing Integers

問題概要

$N$ 頂点の木に1から $N$ までの数字を昇順に1つずつ書き込む。
$v$ に1を書き込み、その後は記入済みの頂点に隣接する頂点に書き込んでいく。
$v=1, 2, \dots, N$ それぞれについて書き込み方の総数（を $10^9+7$ で割った余り）を求めよ。

解法

$cnt_v$ を $v$ を根とする部分木に $v$ から数字を書き始める書き込み方の総数、 $size_v$ を $v$ の子孫の数（ $v$ 自身も含む）とします。
$size_v =\sum_{c \in ch(v)}size_c +1$ は簡単にわかります。
$cnt_v$ は、2から $size_v$ の数字をどの子部分木に割り当てるか決めた後、条件を満たすように数字を書き込む場合の数であるため、
$\begin{eqnarray} cnt_v &=& \frac{(size_v-1)!}{ \prod_{c \in ch(v)}size_c!} \times \prod_{c \in ch(v)}cnt_c \\ &=& (size_v-1)! \prod_{c \in ch(v)} \frac{cnt_c}{size_c!} \end{eqnarray}$

となります。
数えるものが2つあるのでテンプレが使えないように思えますが、 $cnt_v$ と $size_v$ をペアで考える、すなわち $T = \mathbb{N} \times \mathbb{N}$ とすることで上手くいきます。
求める関数は（ $f, g$ の第2引数は使わないので省略）