Dynamic programming on a tree with rerooting technique (tools/rerooting_dp.hpp)

View this file on GitHub
View document part on GitHub
Last update: 2025-01-25 11:01:20+09:00
Include: #include "tools/rerooting_dp.hpp"

It is an abstract framework for dynamic programming on a tree with rerooting technique.

Example

It calculates the diameter of a given tree.

struct monoid {
  using T = int;
  static T op(const T x, const T y) {
    return std::max(x, y);
  }
  static T e() {
    return 0;
  }
};

int main() {
  int n;
  std::cin >> n;

  std::vector<int> costs;
  const auto f_ve = [&](const int dist, const std::size_t edge_id) {
    return dist + costs[edge_id];
  };
  const auto f_ev = [](const int dist, std::size_t) {
    return dist;
  };

  tools::rerooting_dp<int, monoid, decltype(f_ve), decltype(f_ev)> graph(n, f_ve, f_ev);
  for (int i = 0; i < n - 1; ++i) {
    int s, t, w;
    std::cin >> s >> t >> w;
    graph.add_edge(s, t);
    costs.push_back(w);
  }

  const std::vector<int> result = graph.query();
  std::cout << *std::max_element(result.begin(), result.end()) << '\n';
  return 0;
}

License

Author

anqooqie

Constructor

rerooting_dp<R, M, F_VE, F_EV> graph(std::size_t n, F_VE f_ve, F_EV f_ev);

It creates a graph with $n$ vertices and $0$ edges. The meaning of each the type parameter is as follows.

R
- the type of aggregated vertex attributes representing the vertex and the subtree far from that
M
- the commutative monoid on the type of aggregated edge attributes representing the edge and the subtree far from that
f_ve
- the function which accepts an aggregated vertex attribute and the index of the edge adjacent to the subtree, and returns the aggregated edge attribute
f_ev
- the function which accepts the product of aggregated edge attributes and the index of the vertex adjacent to the subtrees, and returns the aggregated vertex attribute

Constraints

For all $a$ in typename M::T and $b$ in typename M::T, M::op(a, b) $=$ M::op(b, a).
For all $a$ in typename M::T, $b$ in typename M::T and $c$ in typename M::T, M::op(M::op(a, b), c) $=$ M::op(a, M::op(b, c)).
For all $a$ in typename M::T, M::op(M::e(), a) $=$ M::op(a, M::e()) $=$ a.
f_ve is invocable.
f_ve accepts two arguments, R and std::size_t.
f_ve returns typename M::T.
f_ev is invocable.
f_ev accepts two arguments, typename M::T and std::size_t.
f_ev returns R.

Time Complexity

$O(1)$

size

std::size_t graph.size();

It returns $n$.

Constraints

None

Time Complexity

$O(1)$

add_edge

std::size_t graph.add_edge(std::size_t u, std::size_t v);

It adds an undirected edge between $u$ and $v$ to the graph. It returns an integer $k$ such that this is the $k$-th edge that is added.

Constraints

$0 \leq u < n$
$0 \leq v < n$

Time Complexity

$O(1)$ amortized

query

std::vector<R> graph.query();

It returns aggregated vertex attributes. The $i$-th element of the return value represents the rooted tree whose root is the vertex $i$.

Constraints

The graph is a tree.

Time Complexity

$O(n)$

Verified with

Code

#ifndef TOOLS_REROOTING_DP_HPP
#define TOOLS_REROOTING_DP_HPP

#include <cstddef>
#include <vector>
#include <cassert>
#include <limits>
#include <stack>
#include <tuple>
#include <algorithm>

namespace tools {
  template <typename R, typename M, typename F_VE, typename F_EV>
  class rerooting_dp {
  private:
    ::std::vector<::std::size_t> m_edges;
    ::std::vector<::std::vector<::std::size_t>> m_graph;
    F_VE m_f_ve;
    F_EV m_f_ev;

    class vertex {
    private:
      const ::tools::rerooting_dp<R, M, F_VE, F_EV> *m_self;

    public:
      ::std::size_t id;
      ::std::size_t neighbor_id_of_parent;
      ::std::vector<::std::size_t> neighbor_ids_of_children;
      typename M::T parent_dp;
      ::std::vector<typename M::T> children_dp;
      ::std::vector<typename M::T> children_dp_cumsum1;
      ::std::vector<typename M::T> children_dp_cumsum2;

      vertex() = default;
      vertex(const vertex&) = default;
      vertex(vertex&&) = default;
      ~vertex() = default;
      vertex& operator=(const vertex&) = default;
      vertex& operator=(vertex&&) = default;

      explicit vertex(const ::tools::rerooting_dp<R, M, F_VE, F_EV> * const self, const ::std::size_t id) :
        m_self(self), id(id), parent_dp(M::e()) {
      }

      ::std::size_t parent_edge_id() const {
        return this->m_self->m_graph[this->id][this->neighbor_id_of_parent];
      }
      ::std::size_t parent_vertex_id() const {
        return this->m_self->m_edges[this->parent_edge_id()] ^ this->id;
      }
      ::std::size_t child_size() const {
        return this->neighbor_ids_of_children.size();
      }
      ::std::size_t child_edge_id(const ::std::size_t child_number) const {
        return this->m_self->m_graph[this->id][this->neighbor_ids_of_children[child_number]];
      }
      ::std::size_t child_vertex_id(const ::std::size_t child_number) const {
        return this->m_self->m_edges[this->child_edge_id(child_number)] ^ this->id;
      }
      R dp_as_root() const {
        return this->m_self->m_f_ev(M::op(this->parent_dp, this->children_dp_cumsum1.back()), this->id);
      }
      R dp_excluding_parent() const {
        return this->m_self->m_f_ev(this->children_dp_cumsum1.back(), this->id);
      }
      R dp_excluding_child(const ::std::size_t excluded_child_number) const {
        return this->m_self->m_f_ev(M::op(this->parent_dp, M::op(this->children_dp_cumsum1[excluded_child_number], this->children_dp_cumsum2[excluded_child_number + 1])), this->id);
      }
    };

  public:
    rerooting_dp() = default;
    rerooting_dp(const ::tools::rerooting_dp<R, M, F_VE, F_EV>&) = default;
    rerooting_dp(::tools::rerooting_dp<R, M, F_VE, F_EV>&&) = default;
    ~rerooting_dp() = default;
    ::tools::rerooting_dp<R, M, F_VE, F_EV>& operator=(const ::tools::rerooting_dp<R, M, F_VE, F_EV>&) = default;
    ::tools::rerooting_dp<R, M, F_VE, F_EV>& operator=(::tools::rerooting_dp<R, M, F_VE, F_EV>&&) = default;

    rerooting_dp(const ::std::size_t n, const F_VE& f_ve, const F_EV& f_ev) : m_graph(n), m_f_ve(f_ve), m_f_ev(f_ev) {
      assert(n >= 1);
    }

    ::std::size_t size() const {
      return this->m_graph.size();
    }

    ::std::size_t add_edge(const ::std::size_t u, const ::std::size_t v) {
      this->m_graph[u].push_back(this->m_edges.size());
      this->m_graph[v].push_back(this->m_edges.size());
      this->m_edges.push_back(u ^ v);
      return this->m_edges.size() - 1;
    }

    ::std::vector<R> query() const {
      assert(this->m_edges.size() + 1 == this->size());

      const int PRE_VERTEX = 1;
      const int POST_EDGE = 2;
      const int POST_VERTEX = 3;
      const ::std::size_t INVALID = ::std::numeric_limits<::std::size_t>::max();

      ::std::vector<vertex> vertices;
      for (::std::size_t i = 0; i < this->size(); ++i) {
        vertices.emplace_back(this, i);
      }

      ::std::stack<::std::tuple<int, ::std::size_t, ::std::size_t>> stack;
      ::std::vector<bool> will_visit(this->size(), false);
      stack.emplace(PRE_VERTEX, 0, INVALID);
      will_visit[0] = true;
      while (!stack.empty()) {
        const int type = ::std::get<0>(stack.top());
        if (type == PRE_VERTEX) {

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];
          stack.emplace(POST_VERTEX, vertex_id, INVALID);
          for (::std::size_t neighbor_id = 0; neighbor_id < this->m_graph[vertex_id].size(); ++neighbor_id) {
            const ::std::size_t child_vertex_id = this->m_edges[this->m_graph[vertex_id][neighbor_id]] ^ vertex_id;
            if (will_visit[child_vertex_id]) {
              v.neighbor_id_of_parent = neighbor_id;
            } else {
              v.neighbor_ids_of_children.push_back(neighbor_id);
              stack.emplace(POST_EDGE, vertex_id, v.child_size() - 1);
              stack.emplace(PRE_VERTEX, child_vertex_id, INVALID);
              will_visit[child_vertex_id] = true;
            }
          }
          v.children_dp.resize(v.child_size());

        } else if (type == POST_EDGE) {

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          const ::std::size_t child_number = ::std::get<2>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];
          const vertex& c = vertices[v.child_vertex_id(child_number)];
          v.children_dp[child_number] = this->m_f_ve(c.dp_excluding_parent(), v.child_edge_id(child_number));

        } else { // POST_VERTEX

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];

          v.children_dp_cumsum1.reserve(v.child_size() + 1);
          v.children_dp_cumsum1.push_back(M::e());
          for (::std::size_t child_number = 0; child_number < v.child_size(); ++child_number) {
            v.children_dp_cumsum1.push_back(M::op(v.children_dp_cumsum1.back(), v.children_dp[child_number]));
          }

          v.children_dp_cumsum2.reserve(v.child_size() + 1);
          v.children_dp_cumsum2.push_back(M::e());
          for (::std::size_t child_number = v.child_size(); child_number --> 0;) {
            v.children_dp_cumsum2.push_back(M::op(v.children_dp[child_number], v.children_dp_cumsum2.back()));
          }
          ::std::reverse(v.children_dp_cumsum2.begin(), v.children_dp_cumsum2.end());

        }
      }

      stack.emplace(PRE_VERTEX, 0, INVALID);
      while (!stack.empty()) {
        const ::std::size_t vertex_id = ::std::get<1>(stack.top());
        stack.pop();

        const vertex& v = vertices[vertex_id];
        for (::std::size_t child_number = 0; child_number < v.child_size(); ++child_number) {
          vertex& c = vertices[v.child_vertex_id(child_number)];
          c.parent_dp = this->m_f_ve(v.dp_excluding_child(child_number), c.parent_edge_id());
          stack.emplace(PRE_VERTEX, c.id, INVALID);
        }
      }

      ::std::vector<R> result;
      result.reserve(this->size());
      for (const vertex& v : vertices) {
        result.push_back(v.dp_as_root());
      }
      return result;
    }
  };
}

#endif

#line 1 "tools/rerooting_dp.hpp"



#include <cstddef>
#include <vector>
#include <cassert>
#include <limits>
#include <stack>
#include <tuple>
#include <algorithm>

namespace tools {
  template <typename R, typename M, typename F_VE, typename F_EV>
  class rerooting_dp {
  private:
    ::std::vector<::std::size_t> m_edges;
    ::std::vector<::std::vector<::std::size_t>> m_graph;
    F_VE m_f_ve;
    F_EV m_f_ev;

    class vertex {
    private:
      const ::tools::rerooting_dp<R, M, F_VE, F_EV> *m_self;

    public:
      ::std::size_t id;
      ::std::size_t neighbor_id_of_parent;
      ::std::vector<::std::size_t> neighbor_ids_of_children;
      typename M::T parent_dp;
      ::std::vector<typename M::T> children_dp;
      ::std::vector<typename M::T> children_dp_cumsum1;
      ::std::vector<typename M::T> children_dp_cumsum2;

      vertex() = default;
      vertex(const vertex&) = default;
      vertex(vertex&&) = default;
      ~vertex() = default;
      vertex& operator=(const vertex&) = default;
      vertex& operator=(vertex&&) = default;

      explicit vertex(const ::tools::rerooting_dp<R, M, F_VE, F_EV> * const self, const ::std::size_t id) :
        m_self(self), id(id), parent_dp(M::e()) {
      }

      ::std::size_t parent_edge_id() const {
        return this->m_self->m_graph[this->id][this->neighbor_id_of_parent];
      }
      ::std::size_t parent_vertex_id() const {
        return this->m_self->m_edges[this->parent_edge_id()] ^ this->id;
      }
      ::std::size_t child_size() const {
        return this->neighbor_ids_of_children.size();
      }
      ::std::size_t child_edge_id(const ::std::size_t child_number) const {
        return this->m_self->m_graph[this->id][this->neighbor_ids_of_children[child_number]];
      }
      ::std::size_t child_vertex_id(const ::std::size_t child_number) const {
        return this->m_self->m_edges[this->child_edge_id(child_number)] ^ this->id;
      }
      R dp_as_root() const {
        return this->m_self->m_f_ev(M::op(this->parent_dp, this->children_dp_cumsum1.back()), this->id);
      }
      R dp_excluding_parent() const {
        return this->m_self->m_f_ev(this->children_dp_cumsum1.back(), this->id);
      }
      R dp_excluding_child(const ::std::size_t excluded_child_number) const {
        return this->m_self->m_f_ev(M::op(this->parent_dp, M::op(this->children_dp_cumsum1[excluded_child_number], this->children_dp_cumsum2[excluded_child_number + 1])), this->id);
      }
    };

  public:
    rerooting_dp() = default;
    rerooting_dp(const ::tools::rerooting_dp<R, M, F_VE, F_EV>&) = default;
    rerooting_dp(::tools::rerooting_dp<R, M, F_VE, F_EV>&&) = default;
    ~rerooting_dp() = default;
    ::tools::rerooting_dp<R, M, F_VE, F_EV>& operator=(const ::tools::rerooting_dp<R, M, F_VE, F_EV>&) = default;
    ::tools::rerooting_dp<R, M, F_VE, F_EV>& operator=(::tools::rerooting_dp<R, M, F_VE, F_EV>&&) = default;

    rerooting_dp(const ::std::size_t n, const F_VE& f_ve, const F_EV& f_ev) : m_graph(n), m_f_ve(f_ve), m_f_ev(f_ev) {
      assert(n >= 1);
    }

    ::std::size_t size() const {
      return this->m_graph.size();
    }

    ::std::size_t add_edge(const ::std::size_t u, const ::std::size_t v) {
      this->m_graph[u].push_back(this->m_edges.size());
      this->m_graph[v].push_back(this->m_edges.size());
      this->m_edges.push_back(u ^ v);
      return this->m_edges.size() - 1;
    }

    ::std::vector<R> query() const {
      assert(this->m_edges.size() + 1 == this->size());

      const int PRE_VERTEX = 1;
      const int POST_EDGE = 2;
      const int POST_VERTEX = 3;
      const ::std::size_t INVALID = ::std::numeric_limits<::std::size_t>::max();

      ::std::vector<vertex> vertices;
      for (::std::size_t i = 0; i < this->size(); ++i) {
        vertices.emplace_back(this, i);
      }

      ::std::stack<::std::tuple<int, ::std::size_t, ::std::size_t>> stack;
      ::std::vector<bool> will_visit(this->size(), false);
      stack.emplace(PRE_VERTEX, 0, INVALID);
      will_visit[0] = true;
      while (!stack.empty()) {
        const int type = ::std::get<0>(stack.top());
        if (type == PRE_VERTEX) {

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];
          stack.emplace(POST_VERTEX, vertex_id, INVALID);
          for (::std::size_t neighbor_id = 0; neighbor_id < this->m_graph[vertex_id].size(); ++neighbor_id) {
            const ::std::size_t child_vertex_id = this->m_edges[this->m_graph[vertex_id][neighbor_id]] ^ vertex_id;
            if (will_visit[child_vertex_id]) {
              v.neighbor_id_of_parent = neighbor_id;
            } else {
              v.neighbor_ids_of_children.push_back(neighbor_id);
              stack.emplace(POST_EDGE, vertex_id, v.child_size() - 1);
              stack.emplace(PRE_VERTEX, child_vertex_id, INVALID);
              will_visit[child_vertex_id] = true;
            }
          }
          v.children_dp.resize(v.child_size());

        } else if (type == POST_EDGE) {

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          const ::std::size_t child_number = ::std::get<2>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];
          const vertex& c = vertices[v.child_vertex_id(child_number)];
          v.children_dp[child_number] = this->m_f_ve(c.dp_excluding_parent(), v.child_edge_id(child_number));

        } else { // POST_VERTEX

          const ::std::size_t vertex_id = ::std::get<1>(stack.top());
          stack.pop();

          vertex& v = vertices[vertex_id];

          v.children_dp_cumsum1.reserve(v.child_size() + 1);
          v.children_dp_cumsum1.push_back(M::e());
          for (::std::size_t child_number = 0; child_number < v.child_size(); ++child_number) {
            v.children_dp_cumsum1.push_back(M::op(v.children_dp_cumsum1.back(), v.children_dp[child_number]));
          }

          v.children_dp_cumsum2.reserve(v.child_size() + 1);
          v.children_dp_cumsum2.push_back(M::e());
          for (::std::size_t child_number = v.child_size(); child_number --> 0;) {
            v.children_dp_cumsum2.push_back(M::op(v.children_dp[child_number], v.children_dp_cumsum2.back()));
          }
          ::std::reverse(v.children_dp_cumsum2.begin(), v.children_dp_cumsum2.end());

        }
      }

      stack.emplace(PRE_VERTEX, 0, INVALID);
      while (!stack.empty()) {
        const ::std::size_t vertex_id = ::std::get<1>(stack.top());
        stack.pop();

        const vertex& v = vertices[vertex_id];
        for (::std::size_t child_number = 0; child_number < v.child_size(); ++child_number) {
          vertex& c = vertices[v.child_vertex_id(child_number)];
          c.parent_dp = this->m_f_ve(v.dp_excluding_child(child_number), c.parent_edge_id());
          stack.emplace(PRE_VERTEX, c.id, INVALID);
        }
      }

      ::std::vector<R> result;
      result.reserve(this->size());
      for (const vertex& v : vertices) {
        result.push_back(v.dp_as_root());
      }
      return result;
    }
  };
}