rerooting_dp.hpp

SECTIONGraph INCLUDEnoya/rerooting_dp.hpp

Generic rerooting DP on a tree. merge combines neighbor contributions, add_edge(value, from, to) moves a rooted value across an edge, and add_vertex(value, vertex) closes a vertex.

在线性时间内计算以每个顶点为根的树 DP；适合答案由相邻子树贡献可结合得到的全根查询。

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_REROOTING_DP_HPP
#define NOYA_REROOTING_DP_HPP 1

/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n).

#include <cassert>
#include <vector>

namespace noya {

/// @brief Generic rerooting DP on a tree.
/// merge combines neighbor contributions, add_edge(value, from, to) moves a
/// rooted value across an edge, and add_vertex(value, vertex) closes a vertex.
template <class T, class Merge, class AddEdge, class AddVertex>
std::vector<T> rerooting_dp(const std::vector<std::vector<int>> &tree,
                            const T &unit, Merge merge, AddEdge add_edge,
                            AddVertex add_vertex, int root = 0) {
  const int n = int(tree.size());
  if (n == 0) {
    return {};
  }
  assert(0 <= root && root < n);
  int degree_sum = 0;
  for (const auto &neighbors : tree) {
    degree_sum += int(neighbors.size());
  }
  assert(degree_sum == 2 * (n - 1));

  std::vector<int> parent(n, -1);
  std::vector<int> order = {root};
  parent[root] = root;
  for (int i = 0; i < int(order.size()); i++) {
    int u = order[i];
    for (int v : tree[u]) {
      assert(0 <= v && v < n);
      if (parent[v] == -1) {
        parent[v] = u;
        order.push_back(v);
      }
    }
  }
  assert(int(order.size()) == n);

  std::vector<T> down(n, unit);
  for (int i = n - 1; i >= 0; i--) {
    int u = order[i];
    T value = unit;
    for (int v : tree[u]) {
      if (parent[v] == u) {
        value = merge(value, add_edge(down[v], v, u));
      }
    }
    down[u] = add_vertex(value, u);
  }

  std::vector<T> from_parent(n, unit);
  std::vector<T> answer(n, unit);
  for (int u : order) {
    int degree = int(tree[u].size());
    std::vector<T> contribution(degree, unit);
    for (int i = 0; i < degree; i++) {
      int v = tree[u][i];
      contribution[i] =
          v == parent[u] ? from_parent[u] : add_edge(down[v], v, u);
    }
    std::vector<T> prefix(degree + 1, unit);
    std::vector<T> suffix(degree + 1, unit);
    for (int i = 0; i < degree; i++) {
      prefix[i + 1] = merge(prefix[i], contribution[i]);
    }
    for (int i = degree - 1; i >= 0; i--) {
      suffix[i] = merge(contribution[i], suffix[i + 1]);
    }
    answer[u] = add_vertex(prefix[degree], u);
    for (int i = 0; i < degree; i++) {
      int v = tree[u][i];
      if (parent[v] == u) {
        T without_child = merge(prefix[i], suffix[i + 1]);
        from_parent[v] = add_edge(add_vertex(without_child, u), u, v);
      }
    }
  }
  return answer;
}

} // namespace noya

#endif // NOYA_REROOTING_DP_HPP

#include <cassert>
#include <vector>

/// @complexity Time: O(n) monoid transitions on a tree.
/// Space: O(n).

namespace noya {

/// @brief Generic rerooting DP on a tree.
/// merge combines neighbor contributions, add_edge(value, from, to) moves a
/// rooted value across an edge, and add_vertex(value, vertex) closes a vertex.
template <class T, class Merge, class AddEdge, class AddVertex>
std::vector<T> rerooting_dp(const std::vector<std::vector<int>> &tree,
                            const T &unit, Merge merge, AddEdge add_edge,
                            AddVertex add_vertex, int root = 0) {
  const int n = int(tree.size());
  if (n == 0) {
    return {};
  }
  assert(0 <= root && root < n);
  int degree_sum = 0;
  for (const auto &neighbors : tree) {
    degree_sum += int(neighbors.size());
  }
  assert(degree_sum == 2 * (n - 1));

  std::vector<int> parent(n, -1);
  std::vector<int> order = {root};
  parent[root] = root;
  for (int i = 0; i < int(order.size()); i++) {
    int u = order[i];
    for (int v : tree[u]) {
      assert(0 <= v && v < n);
      if (parent[v] == -1) {
        parent[v] = u;
        order.push_back(v);
      }
    }
  }
  assert(int(order.size()) == n);

  std::vector<T> down(n, unit);
  for (int i = n - 1; i >= 0; i--) {
    int u = order[i];
    T value = unit;
    for (int v : tree[u]) {
      if (parent[v] == u) {
        value = merge(value, add_edge(down[v], v, u));
      }
    }
    down[u] = add_vertex(value, u);
  }

  std::vector<T> from_parent(n, unit);
  std::vector<T> answer(n, unit);
  for (int u : order) {
    int degree = int(tree[u].size());
    std::vector<T> contribution(degree, unit);
    for (int i = 0; i < degree; i++) {
      int v = tree[u][i];
      contribution[i] =
          v == parent[u] ? from_parent[u] : add_edge(down[v], v, u);
    }
    std::vector<T> prefix(degree + 1, unit);
    std::vector<T> suffix(degree + 1, unit);
    for (int i = 0; i < degree; i++) {
      prefix[i + 1] = merge(prefix[i], contribution[i]);
    }
    for (int i = degree - 1; i >= 0; i--) {
      suffix[i] = merge(contribution[i], suffix[i + 1]);
    }
    answer[u] = add_vertex(prefix[degree], u);
    for (int i = 0; i < degree; i++) {
      int v = tree[u][i];
      if (parent[v] == u) {
        T without_child = merge(prefix[i], suffix[i + 1]);
        from_parent[v] = add_edge(add_vertex(without_child, u), u, v);
      }
    }
  }
  return answer;
}

} // namespace noya