全方位木 DP
(tree/rerooting-dp.hpp)

View this file on GitHub
Last update: 2026-07-15 18:17:55+09:00
Include: #include "tree/rerooting-dp.hpp"
Link: View error logs on GitHub Actions

全方位木 DP を行う．

各頂点を根としたときの木 DP の値をすべて求める．森にも使える． DP の型 T は可換モノイドであることを仮定する．

ReRootingDP<T, V, E, merge, add_edge, add_vertex>(n, id)：$n$ 頂点で初期化する．
- id：T の単位元．
- merge(x, y)：部分木 DP の併合．
- add_edge(x, e)：DP 値 x に辺データ e を追加する．
- add_vertex(x, v)：DP 値 x に頂点データ v を追加する．
set_vertex(u, data)：頂点 $u$ のデータを設定する．
add_edge(u, v, data)：無向辺 $uv$ を追加し，辺 ID を返す．
add_edge(u, v, data, rdata)：$u \to v$ 方向の辺データを data，$v \to u$ 方向の辺データを rdata として無向辺を追加し，辺 ID を返す．
calc()：DP を計算し，各頂点を根とした答えを返す．
dp[i]：calc() 後，頂点 $i$ を根とした答えを返す．
dp.edge_subtree(eid, side)：calc() 後，辺 eid を切った片側の DP 値を返す．side = 0 なら add_edge の第 1 引数側，side = 1 なら第 2 引数側．

時間計算量は，merge，add_edge，add_vertex が $O(1)$ のとき calc() が $O(N)$，edge_subtree が $O(1)$．

Verified with

verify/tree/LC_tree_path_composite_sum.test.cpp

Code

#pragma once

template <class T, class V, class E, T (*merge)(T, T), T (*add_edge_fn)(T, E), T (*add_vertex_fn)(T, V)>
struct ReRootingDP {
 private:
  struct Edge {
    int to, idx, rev;
    E data;
  };

  int n;
  T id;
  vector<vector<Edge>> g;
  vector<V> vertex_data;
  vector<T> dp;
  vector<T> ans;
  vector<array<int, 2>> edge_pos;
  int edge_count = 0;

 public:
  ReRootingDP() {}
  ReRootingDP(int n, T id) : n(n), id(id), g(n), vertex_data(n) {}
  void set_vertex(int u, V data) { vertex_data[u] = data; }
  int add_edge(int u, int v, E data) { return add_edge(u, v, data, data); }
  int add_edge(int u, int v, E data, E rdata) {
    int did = dp.size();
    dp.push_back(id);
    dp.push_back(id);
    g[u].push_back({v, did, did + 1, data});
    g[v].push_back({u, did + 1, did, rdata});
    edge_pos.push_back({did, did + 1});
    return edge_count++;
  }
  vector<T> calc() {
    fill(dp.begin(), dp.end(), id);
    ans.assign(n, id);
    int max_deg = 0;
    for (int i = 0; i < n; i++) max_deg = max(max_deg, (int)g[i].size());
    vector<int> parent(n, -2), parent_edge(n, -1), order;
    order.reserve(n);
    for (int root = 0; root < n; root++) {
      if (parent[root] != -2) continue;
      parent[root] = -1;
      stack<int> st;
      st.push(root);
      while (!st.empty()) {
        int x = st.top();
        st.pop();
        order.push_back(x);
        for (int i = 0; i < (int)g[x].size(); i++) {
          int y = g[x][i].to;
          if (parent[y] != -2) continue;
          parent[y] = x;
          parent_edge[y] = g[x][i].rev;
          st.push(y);
        }
      }
    }
    for (int i = (int)order.size() - 1; i >= 0; i--) {
      int x = order[i];
      T val = id;
      for (auto& e : g[x]) {
        if (e.to == parent[x]) continue;
        val = merge(val, add_edge_fn(dp[e.rev], e.data));
      }
      if (parent[x] == -1)
        ans[x] = add_vertex_fn(val, vertex_data[x]);
      else
        dp[parent_edge[x]] = add_vertex_fn(val, vertex_data[x]);
    }
    vector<T> sum_r(max_deg + 1, id);
    for (int x : order) {
      int deg = g[x].size();
      sum_r[deg] = id;
      for (int i = deg - 1; i >= 0; i--) {
        auto& e = g[x][i];
        sum_r[i] = merge(sum_r[i + 1], add_edge_fn(dp[e.rev], e.data));
      }
      ans[x] = add_vertex_fn(sum_r[0], vertex_data[x]);
      T sum_l = id;
      for (int i = 0; i < deg; i++) {
        auto& e = g[x][i];
        if (e.to != parent[x]) {
          dp[e.idx] = add_vertex_fn(merge(sum_l, sum_r[i + 1]), vertex_data[x]);
        }
        sum_l = merge(sum_l, add_edge_fn(dp[e.rev], e.data));
      }
    }
    return ans;
  }
  T operator[](int i) const { return ans[i]; }
  T edge_subtree(int edge_id, int side) const {
    assert(0 <= edge_id && edge_id < edge_count);
    assert(0 <= side && side < 2);
    return dp[edge_pos[edge_id][side]];
  }
};

/**
 * @brief 全方位木 DP
 * @docs docs/tree/rerooting-dp.md
 */

#line 2 "tree/rerooting-dp.hpp"

template <class T, class V, class E, T (*merge)(T, T), T (*add_edge_fn)(T, E), T (*add_vertex_fn)(T, V)>
struct ReRootingDP {
 private:
  struct Edge {
    int to, idx, rev;
    E data;
  };

  int n;
  T id;
  vector<vector<Edge>> g;
  vector<V> vertex_data;
  vector<T> dp;
  vector<T> ans;
  vector<array<int, 2>> edge_pos;
  int edge_count = 0;

 public:
  ReRootingDP() {}
  ReRootingDP(int n, T id) : n(n), id(id), g(n), vertex_data(n) {}
  void set_vertex(int u, V data) { vertex_data[u] = data; }
  int add_edge(int u, int v, E data) { return add_edge(u, v, data, data); }
  int add_edge(int u, int v, E data, E rdata) {
    int did = dp.size();
    dp.push_back(id);
    dp.push_back(id);
    g[u].push_back({v, did, did + 1, data});
    g[v].push_back({u, did + 1, did, rdata});
    edge_pos.push_back({did, did + 1});
    return edge_count++;
  }
  vector<T> calc() {
    fill(dp.begin(), dp.end(), id);
    ans.assign(n, id);
    int max_deg = 0;
    for (int i = 0; i < n; i++) max_deg = max(max_deg, (int)g[i].size());
    vector<int> parent(n, -2), parent_edge(n, -1), order;
    order.reserve(n);
    for (int root = 0; root < n; root++) {
      if (parent[root] != -2) continue;
      parent[root] = -1;
      stack<int> st;
      st.push(root);
      while (!st.empty()) {
        int x = st.top();
        st.pop();
        order.push_back(x);
        for (int i = 0; i < (int)g[x].size(); i++) {
          int y = g[x][i].to;
          if (parent[y] != -2) continue;
          parent[y] = x;
          parent_edge[y] = g[x][i].rev;
          st.push(y);
        }
      }
    }
    for (int i = (int)order.size() - 1; i >= 0; i--) {
      int x = order[i];
      T val = id;
      for (auto& e : g[x]) {
        if (e.to == parent[x]) continue;
        val = merge(val, add_edge_fn(dp[e.rev], e.data));
      }
      if (parent[x] == -1)
        ans[x] = add_vertex_fn(val, vertex_data[x]);
      else
        dp[parent_edge[x]] = add_vertex_fn(val, vertex_data[x]);
    }
    vector<T> sum_r(max_deg + 1, id);
    for (int x : order) {
      int deg = g[x].size();
      sum_r[deg] = id;
      for (int i = deg - 1; i >= 0; i--) {
        auto& e = g[x][i];
        sum_r[i] = merge(sum_r[i + 1], add_edge_fn(dp[e.rev], e.data));
      }
      ans[x] = add_vertex_fn(sum_r[0], vertex_data[x]);
      T sum_l = id;
      for (int i = 0; i < deg; i++) {
        auto& e = g[x][i];
        if (e.to != parent[x]) {
          dp[e.idx] = add_vertex_fn(merge(sum_l, sum_r[i + 1]), vertex_data[x]);
        }
        sum_l = merge(sum_l, add_edge_fn(dp[e.rev], e.data));
      }
    }
    return ans;
  }
  T operator[](int i) const { return ans[i]; }
  T edge_subtree(int edge_id, int side) const {
    assert(0 <= edge_id && edge_id < edge_count);
    assert(0 <= side && side < 2);
    return dp[edge_pos[edge_id][side]];
  }
};

/**
 * @brief 全方位木 DP
 * @docs docs/tree/rerooting-dp.md
 */

全方位木 DP (tree/rerooting-dp.hpp)

Verified with

Code

全方位木 DP
(tree/rerooting-dp.hpp)