Library/Graph/Tree/ReRootingDP.hpp

• View this file on GitHub
• Last update: 2025-06-14 20:53:47+09:00
• Include: #include "Library/Graph/Tree/ReRootingDP.hpp"

Depends on

• Library/Algebraic/Monoid.hpp
• Library/Graph/Graph.hpp
• Library/Graph/Normal/BFS.hpp
• Library/Graph/Tree/TreeDP.hpp

Verified with

• Test/Graph/Tree/ReRootingDP.test.cpp
• Test/Graph/Tree/ReRootingDP_cost.test.cpp

Code

#pragma once
#include <unordered_map>

#include <vector>


#include "../../Algebraic/Monoid.hpp"

#include "../../Graph/Normal/BFS.hpp"

#include "../../Graph/Tree/TreeDP.hpp"


namespace mtd {
  /*
   * Monoid: 部分木の情報
   * edge_f: 辺の情報を親に流す関数: (M, f, t, c) -> M
   * node_f: 子の情報を親に流す関数: (M, i) -> M
   */
  template <monoid Monoid, class Node, class Cost, class Lambda1, class Lambda2>
  auto reRootingDP(const Graph<Node, Cost>& graph, const Lambda1& edge_f,
                   const Lambda2& node_f) {
    constexpr int root = 0;
    auto n = graph.size();

    // <辺情報を考慮したMonoidの2項演算>

    auto merge = [&](Monoid& m1, const Monoid& m2, Node f = -1, Node t = -1,
                     const Cost& c = Cost()) {
      m1 = m1.binaryOperation((f == -1 || t == -1) ? m2 : edge_f(m2, f, t, c));
    };

    // <node:toを根とした木で全てマージした解を求める>

    std::vector<std::vector<std::tuple<Monoid, Node, Cost>>> partial(n);
    auto all_merge = [&](Node to) {
      Monoid val{};
      for (const auto& [ad, from, cost] : partial[to]) {
        merge(val, ad, from, to, cost);
      }
      return node_f(val, to);
    };

    // <node:toを根とした木でfrom以外マージした解を求める>

    std::vector<std::unordered_map<Node, Monoid>> partial_ac(n);
    std::vector<Monoid> ret_m(n);
    auto accumulation = [&](Node to) {
      // 左からマージ

      Monoid val_ord{};
      for (const auto& [ad, from, cost] : partial[to]) {
        partial_ac[to].emplace(from, val_ord);
        merge(val_ord, ad, from, to, cost);
      }
      // 右からマージ

      Monoid val_rord{};
      for (auto rit = partial[to].rbegin(); rit != partial[to].rend(); ++rit) {
        auto [ad, from, cost] = *rit;
        merge(partial_ac[to][from], val_rord, cost);
        merge(val_rord, ad, from, to, cost);
      }
      // node情報を反映させて値を確定

      ret_m[to] = node_f(val_ord, to);
      for (auto&& [_, pac] : partial_ac[to]) { pac = node_f(pac, to); }
    };

    // rootを根とした解を求める

    treeDP(graph, root, [&](Node f, Node t, const Cost& c) {
      partial[t].emplace_back(all_merge(f), f, c);
    });
    accumulation(0);

    // rootからbfsして各nodeを根とした解を求める

    bfs(graph, root, [&](Node f, Node t, const Cost& c) {
      partial[t].emplace_back(partial_ac[f][t], f, c);
      accumulation(t);
    });

    std::vector<typename Monoid::value_type> ret;
    for (const auto x : ret_m) { ret.emplace_back(x.m_val); }
    return ret;
  }
}  // namespace mtd

#line 2 "Library/Graph/Tree/ReRootingDP.hpp"
#include <unordered_map>

#include <vector>


#line 2 "Library/Algebraic/Monoid.hpp"

#include <iostream>

namespace mtd {

  template <class S,    // set
            S element,  // identity element
            class op    // binary operation
            >
  requires std::is_invocable_r_v<S, op, S, S>
  struct Monoid {
    using value_type = S;
    constexpr static S _element = element;
    using op_type = op;

    S m_val;
    constexpr Monoid(S val) : m_val(val) {}
    constexpr Monoid() : Monoid(element) {}
    constexpr Monoid binaryOperation(const Monoid& m2) const {
      return op()(m_val, m2.m_val);
    }
    friend std::ostream& operator<<(std::ostream& os,
                                    const Monoid<S, element, op>& m) {
      return os << m.m_val;
    }
  };

  namespace __detail {
    template <typename T, template <typename, auto, typename> typename S>
    concept is_monoid_specialization_of = requires {
      typename std::enable_if_t<std::is_same_v<
          T, S<typename T::value_type, T::_element, typename T::op_type>>>;
    };
  }  // namespace __detail

  template <typename M>
  concept monoid = __detail::is_monoid_specialization_of<M, Monoid>;

}  // namespace mtd
#line 2 "Library/Graph/Normal/BFS.hpp"

#include <concepts>

#include <queue>

#line 6 "Library/Graph/Normal/BFS.hpp"

#line 2 "Library/Graph/Graph.hpp"
#include <deque>

#line 4 "Library/Graph/Graph.hpp"
#include <ranges>

#include <tuple>

#line 7 "Library/Graph/Graph.hpp"

namespace mtd {
  template <class Node = long long, class Cost = long long>
  class Graph {
    using Edge = std::pair<Node, Cost>;
    using Edges = std::vector<Edge>;

    const int m_n;
    std::vector<Edges> m_graph;

  public:
    Graph(int n) : m_n(n), m_graph(n) {}
    Graph(const std::vector<Edges>& edges)
        : m_n(edges.size()), m_graph(edges) {}
    Graph(int n, const std::vector<std::tuple<Node, Node>>& edges,
          bool is_arc = false, bool is_index1 = true)
        : Graph<Node, Cost>(n) {
      for (auto [u, v] : edges) {
        u -= is_index1;
        v -= is_index1;
        if (is_arc) {
          addArc(u, v);
        } else {
          addEdge(u, v);
        }
      }
    }
    Graph(int n, const std::vector<std::tuple<Node, Node, Cost>>& edges,
          bool is_arc = false, bool is_index1 = true)
        : Graph<Node, Cost>(n) {
      for (auto [u, v, c] : edges) {
        u -= is_index1;
        v -= is_index1;
        if (is_arc) {
          addArc(u, v, c);
        } else {
          addEdge(u, v, c);
        }
      }
    }

    auto addEdge(const Node& f, const Node& t, const Cost& c = 1) {
      addArc(f, t, c);
      addArc(t, f, c);
    }
    auto addArc(const Node& f, const Node& t, const Cost& c = 1) {
      m_graph[f].emplace_back(t, c);
    }
    auto getEdges(const Node& from) const {
      class EdgesRange {
        const typename Edges::const_iterator b, e;

      public:
        EdgesRange(const Edges& edges) : b(edges.begin()), e(edges.end()) {}
        auto begin() const { return b; }
        auto end() const { return e; }
      };
      return EdgesRange(m_graph[from]);
    }
    auto getEdges() const {
      std::deque<std::tuple<Node, Node, Cost>> edges;
      for (Node from : std::views::iota(0, m_n)) {
        for (const auto& [to, c] : getEdges(from)) {
          edges.emplace_back(from, to, c);
        }
      }
      return edges;
    }
    auto getEdgesExcludeCost() const {
      std::deque<std::pair<Node, Node>> edges;
      for (Node from : std::views::iota(0, m_n)) {
        for (const auto& [to, _] : getEdges(from)) {
          edges.emplace_back(from, to);
        }
      }
      return edges;
    }
    auto reverse() const {
      auto rev = Graph<Node, Cost>(m_n);
      for (const auto& [from, to, c] : getEdges()) { rev.addArc(to, from, c); }
      return rev;
    }
    auto size() const { return m_n; };
    auto debug(bool directed = false) const {
      for (const auto& [f, t, c] : getEdges()) {
        if (f < t || directed) {
          std::cout << f << " -> " << t << ": " << c << std::endl;
        }
      }
    }
  };
}  // namespace mtd

#line 8 "Library/Graph/Normal/BFS.hpp"

namespace mtd {
  template <class Node, class Cost, class Lambda,
            std::convertible_to<Node> _Node>
  auto bfs(const Graph<Node, Cost>& graph, const _Node& root,
           const Lambda& lambda) {
    auto n = graph.size();
    std::vector<bool> used(n);
    used[root] = true;
    std::queue<Node> q;
    q.emplace(root);
    while (!q.empty()) {
      auto from = q.front();
      q.pop();
      for (const auto& [to, cost] : graph.getEdges(from)) {
        if (used[to]) { continue; }
        q.emplace(to);
        used[to] = true;
        lambda(from, to, cost);
      }
    }
  }
}  // namespace mtd

#line 2 "Library/Graph/Tree/TreeDP.hpp"

#line 6 "Library/Graph/Tree/TreeDP.hpp"

#line 8 "Library/Graph/Tree/TreeDP.hpp"

namespace mtd {
  template <class Node, class Cost, class Lambda,
            std::convertible_to<Node> _Node>
  auto treeDP(const Graph<Node, Cost>& tree, _Node root, const Lambda& lambda) {
    auto n = tree.size();
    std::vector<Node> in(n);
    for (const auto& [f, t] : tree.getEdgesExcludeCost())
      if (f < t) {
        ++in[f];
        ++in[t];
      }
    std::queue<Node> q;
    std::vector<bool> used(n);
    for (Node i = 0; i < n; ++i)
      if (i != root && in[i] == 1) { q.emplace(i); }
    while (!q.empty()) {
      auto from = q.front();
      q.pop();
      used[from] = true;

      for (const auto& [to, cost] : tree.getEdges(from)) {
        if (used[to]) { continue; }
        lambda(from, to, cost);
        --in[to];
        if (to != root && in[to] == 1) { q.emplace(to); }
      }
    }
  }
}  // namespace mtd

#line 8 "Library/Graph/Tree/ReRootingDP.hpp"

namespace mtd {
  /*
   * Monoid: 部分木の情報
   * edge_f: 辺の情報を親に流す関数: (M, f, t, c) -> M
   * node_f: 子の情報を親に流す関数: (M, i) -> M
   */
  template <monoid Monoid, class Node, class Cost, class Lambda1, class Lambda2>
  auto reRootingDP(const Graph<Node, Cost>& graph, const Lambda1& edge_f,
                   const Lambda2& node_f) {
    constexpr int root = 0;
    auto n = graph.size();

    // <辺情報を考慮したMonoidの2項演算>

    auto merge = [&](Monoid& m1, const Monoid& m2, Node f = -1, Node t = -1,
                     const Cost& c = Cost()) {
      m1 = m1.binaryOperation((f == -1 || t == -1) ? m2 : edge_f(m2, f, t, c));
    };

    // <node:toを根とした木で全てマージした解を求める>

    std::vector<std::vector<std::tuple<Monoid, Node, Cost>>> partial(n);
    auto all_merge = [&](Node to) {
      Monoid val{};
      for (const auto& [ad, from, cost] : partial[to]) {
        merge(val, ad, from, to, cost);
      }
      return node_f(val, to);
    };

    // <node:toを根とした木でfrom以外マージした解を求める>

    std::vector<std::unordered_map<Node, Monoid>> partial_ac(n);
    std::vector<Monoid> ret_m(n);
    auto accumulation = [&](Node to) {
      // 左からマージ

      Monoid val_ord{};
      for (const auto& [ad, from, cost] : partial[to]) {
        partial_ac[to].emplace(from, val_ord);
        merge(val_ord, ad, from, to, cost);
      }
      // 右からマージ

      Monoid val_rord{};
      for (auto rit = partial[to].rbegin(); rit != partial[to].rend(); ++rit) {
        auto [ad, from, cost] = *rit;
        merge(partial_ac[to][from], val_rord, cost);
        merge(val_rord, ad, from, to, cost);
      }
      // node情報を反映させて値を確定

      ret_m[to] = node_f(val_ord, to);
      for (auto&& [_, pac] : partial_ac[to]) { pac = node_f(pac, to); }
    };

    // rootを根とした解を求める

    treeDP(graph, root, [&](Node f, Node t, const Cost& c) {
      partial[t].emplace_back(all_merge(f), f, c);
    });
    accumulation(0);

    // rootからbfsして各nodeを根とした解を求める

    bfs(graph, root, [&](Node f, Node t, const Cost& c) {
      partial[t].emplace_back(partial_ac[f][t], f, c);
      accumulation(t);
    });

    std::vector<typename Monoid::value_type> ret;
    for (const auto x : ret_m) { ret.emplace_back(x.m_val); }
    return ret;
  }
}  // namespace mtd

Back to top page