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

• View this file on GitHub
• Last update: 2024-05-03 23:21:26+09:00
• Include: #include "tree/rerooting.hpp"

全方位木DP

テンプレート

// 「T : 根が virtual である根付き木」に対応する情報を管理する
using T = ;
// 空の状態に対応する情報
T leaf = ;
// T 同士をマージ
auto f1 = [&](T x, T y) -> T {
  
};
// T の根に頂点 c および辺 c-p を追加する (p は virtual)
auto f2 = [&](T x, int c, int p) -> T {
  
};
Rerooting<T, decltype(g), decltype(f1), decltype(f2)> dp(g, f1, f2, leaf);

Depends on

• グラフテンプレート (graph/graph-template.hpp)

Verified with

• verify/verify-aoj-grl/aoj-grl-5-a-rerooting.test.cpp
• verify/verify-aoj-grl/aoj-grl-5-b.test.cpp
• verify/verify-unit-test/rerooting.test.cpp
• verify/verify-yosupo-graph/yosupo-tree-path-composite-sum.test.cpp

Code

#pragma once

#include "../graph/graph-template.hpp"

// Rerooting
// f1(c1, c2) ... merge value of child node
// f2(memo[i], chd, par) ... return value from child node to parent node
// memo[i] ... result of subtree rooted i
// dp[i] ... result of tree rooted i
template <typename T, typename G, typename F1, typename F2>
struct Rerooting {
  const G &g;
  const F1 f1;
  const F2 f2;
  vector<T> memo, dp;
  T leaf;

  Rerooting(const G &_g, const F1 _f1, const F2 _f2, const T &_leaf)
      : g(_g), f1(_f1), f2(_f2), memo(g.size()), dp(g.size()), leaf(_leaf) {
    dfs(0, -1);
    dfs2(0, -1, T{});
  }

  const T &operator[](int i) const { return dp[i]; }

  void dfs(int cur, int par) {
    vector<T> chds;
    for (auto &dst : g[cur]) {
      if (dst == par) continue;
      dfs(dst, cur);
      chds.push_back(f2(memo[dst], dst, cur));
    }
    if (chds.empty()) {
      memo[cur] = leaf;
    } else {
      memo[cur] = chds[0];
      for (int i = 1; i < (int)chds.size(); i++) {
        memo[cur] = f1(memo[cur], chds[i]);
      }
    }
  }

  void dfs2(int cur, int par, const T &pval) {
    // get cumulative sum
    vector<T> buf;
    if (cur != 0) buf.push_back(pval);
    for (auto dst : g[cur]) {
      if (dst == par) continue;
      buf.push_back(f2(memo[dst], dst, cur));
    }
    vector<T> head(buf.size()), tail(buf.size());
    if (!buf.empty()) {
      head[0] = buf[0];
      for (int i = 1; i < (int)buf.size(); i++) {
        head[i] = f1(head[i - 1], buf[i]);
      }
      tail.back() = buf.back();
      for (int i = (int)buf.size() - 2; i >= 0; i--) {
        tail[i] = f1(tail[i + 1], buf[i]);
      }
    }
    dp[cur] = head.empty() ? leaf : head.back();
    int idx = cur == 0 ? 0 : 1;
    for (auto &dst : g[cur]) {
      if (dst == par) continue;
      T val;
      bool first = idx == 0;
      bool last = idx + 1 == (int)head.size();
      if (first and last) {
        val = leaf;
      } else if (first) {
        val = tail[idx + 1];
      } else if (last) {
        val = head[idx - 1];
      } else {
        val = f1(head[idx - 1], tail[idx + 1]);
      }
      dfs2(dst, cur, f2(val, cur, dst));
      idx++;
    }
  }
};

/**
 * @brief Rerooting(全方位木DP)
 * @docs docs/tree/rerooting.md
 */

#line 2 "tree/rerooting.hpp"

#line 2 "graph/graph-template.hpp"

template <typename T>
struct edge {
  int src, to;
  T cost;

  edge(int _to, T _cost) : src(-1), to(_to), cost(_cost) {}
  edge(int _src, int _to, T _cost) : src(_src), to(_to), cost(_cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  operator int() const { return to; }
};
template <typename T>
using Edges = vector<edge<T>>;
template <typename T>
using WeightedGraph = vector<Edges<T>>;
using UnweightedGraph = vector<vector<int>>;

// Input of (Unweighted) Graph
UnweightedGraph graph(int N, int M = -1, bool is_directed = false,
                      bool is_1origin = true) {
  UnweightedGraph g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    if (is_1origin) x--, y--;
    g[x].push_back(y);
    if (!is_directed) g[y].push_back(x);
  }
  return g;
}

// Input of Weighted Graph
template <typename T>
WeightedGraph<T> wgraph(int N, int M = -1, bool is_directed = false,
                        bool is_1origin = true) {
  WeightedGraph<T> g(N);
  if (M == -1) M = N - 1;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    cin >> c;
    if (is_1origin) x--, y--;
    g[x].emplace_back(x, y, c);
    if (!is_directed) g[y].emplace_back(y, x, c);
  }
  return g;
}

// Input of Edges
template <typename T>
Edges<T> esgraph([[maybe_unused]] int N, int M, int is_weighted = true,
                 bool is_1origin = true) {
  Edges<T> es;
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    es.emplace_back(x, y, c);
  }
  return es;
}

// Input of Adjacency Matrix
template <typename T>
vector<vector<T>> adjgraph(int N, int M, T INF, int is_weighted = true,
                           bool is_directed = false, bool is_1origin = true) {
  vector<vector<T>> d(N, vector<T>(N, INF));
  for (int _ = 0; _ < M; _++) {
    int x, y;
    cin >> x >> y;
    T c;
    if (is_weighted)
      cin >> c;
    else
      c = 1;
    if (is_1origin) x--, y--;
    d[x][y] = c;
    if (!is_directed) d[y][x] = c;
  }
  return d;
}

/**
 * @brief グラフテンプレート
 * @docs docs/graph/graph-template.md
 */
#line 4 "tree/rerooting.hpp"

// Rerooting
// f1(c1, c2) ... merge value of child node
// f2(memo[i], chd, par) ... return value from child node to parent node
// memo[i] ... result of subtree rooted i
// dp[i] ... result of tree rooted i
template <typename T, typename G, typename F1, typename F2>
struct Rerooting {
  const G &g;
  const F1 f1;
  const F2 f2;
  vector<T> memo, dp;
  T leaf;

  Rerooting(const G &_g, const F1 _f1, const F2 _f2, const T &_leaf)
      : g(_g), f1(_f1), f2(_f2), memo(g.size()), dp(g.size()), leaf(_leaf) {
    dfs(0, -1);
    dfs2(0, -1, T{});
  }

  const T &operator[](int i) const { return dp[i]; }

  void dfs(int cur, int par) {
    vector<T> chds;
    for (auto &dst : g[cur]) {
      if (dst == par) continue;
      dfs(dst, cur);
      chds.push_back(f2(memo[dst], dst, cur));
    }
    if (chds.empty()) {
      memo[cur] = leaf;
    } else {
      memo[cur] = chds[0];
      for (int i = 1; i < (int)chds.size(); i++) {
        memo[cur] = f1(memo[cur], chds[i]);
      }
    }
  }

  void dfs2(int cur, int par, const T &pval) {
    // get cumulative sum
    vector<T> buf;
    if (cur != 0) buf.push_back(pval);
    for (auto dst : g[cur]) {
      if (dst == par) continue;
      buf.push_back(f2(memo[dst], dst, cur));
    }
    vector<T> head(buf.size()), tail(buf.size());
    if (!buf.empty()) {
      head[0] = buf[0];
      for (int i = 1; i < (int)buf.size(); i++) {
        head[i] = f1(head[i - 1], buf[i]);
      }
      tail.back() = buf.back();
      for (int i = (int)buf.size() - 2; i >= 0; i--) {
        tail[i] = f1(tail[i + 1], buf[i]);
      }
    }
    dp[cur] = head.empty() ? leaf : head.back();
    int idx = cur == 0 ? 0 : 1;
    for (auto &dst : g[cur]) {
      if (dst == par) continue;
      T val;
      bool first = idx == 0;
      bool last = idx + 1 == (int)head.size();
      if (first and last) {
        val = leaf;
      } else if (first) {
        val = tail[idx + 1];
      } else if (last) {
        val = head[idx - 1];
      } else {
        val = f1(head[idx - 1], tail[idx + 1]);
      }
      dfs2(dst, cur, f2(val, cur, dst));
      idx++;
    }
  }
};

/**
 * @brief Rerooting(全方位木DP)
 * @docs docs/tree/rerooting.md
 */

Back to top page