Library/Graph/Tree/AuxiliaryTree.hpp

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

Depends on

• Library/Graph/Graph.hpp
• Library/Graph/Tree/HeavyLightDecomposition.hpp

Verified with

• Test/Graph/Tree/AuxiliaryTree.test.cpp

Code

#pragma once

#include <set>

#include <vector>


#include "./HeavyLightDecomposition.hpp"


namespace mtd {
  template <class Node, class Cost>
  class AuxiliaryTree {
    // 定数倍高速化のため破壊的

    std::vector<int> compres_map;

    const std::vector<Cost> depth_cost;
    const HeavyLightDecomposition<Node, Cost> hld;

    auto construct_depth(const Graph<Node, Cost>& tree) const {
      std::vector<Cost> _depth_cost(tree.size());
      std::vector<int> used(tree.size());
      auto dfs = [&](auto&& self, Node from) -> void {
        used[from] = true;
        for (const auto& [to, c] : tree.getEdges(from))
          if (!used[to]) {
            _depth_cost[to] = _depth_cost[from] + c;
            self(self, to);
          }
      };
      dfs(dfs, 0);
      return _depth_cost;
    }

  public:
    AuxiliaryTree(const Graph<Node, Cost>& tree)
        : compres_map(tree.size()),
          depth_cost(construct_depth(tree)),
          hld(tree) {}

    auto compression(const std::vector<Node>& nodes) {
      auto compare = [&](int a, int b) { return hld.getId(a) < hld.getId(b); };

      // 元の頂点集合

      auto nodes_set =
          std::set<int, decltype(compare)>(nodes.begin(), nodes.end(), compare);
      auto nodes_set_with_lca = nodes_set;

      // pre orderでの全ての隣接nodeのLCAを求める

      for (auto itr = nodes_set_with_lca.begin();
           std::next(itr) != nodes_set_with_lca.end(); ++itr) {
        nodes_set_with_lca.emplace(hld.lca(*itr, *std::next(itr)));
      }

      // 座標圧縮

      int at_size = nodes_set_with_lca.size();
      for (int i = 0; auto x : nodes_set_with_lca) {
        compres_map[x] = i;
        ++i;
      }

      // LCAを含めた全てのnodeで子孫関係を保って辺を張る

      std::stack<int> stk;
      Graph<Node, Cost> auxiliary_tree(at_size);
      for (auto nd : nodes_set_with_lca) {
        while (!stk.empty() && hld.lca(stk.top(), nd) != stk.top()) {
          stk.pop();
        }
        if (!stk.empty()) {
          auto f = compres_map[stk.top()];
          auto t = compres_map[nd];
          auto c = depth_cost[stk.top()] + depth_cost[nd] -
                   depth_cost[hld.lca(stk.top(), nd)] * 2;
          auxiliary_tree.addEdge(f, t, c);
        }
        stk.emplace(nd);
      }
      return auxiliary_tree;
    }
  };
}  // namespace mtd

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

#include <set>

#include <vector>


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

#include <queue>

#include <stack>

#include <unordered_map>


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

#include <iostream>

#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/Tree/HeavyLightDecomposition.hpp"

namespace mtd {
  template <class Node, class Cost>
  class HeavyLightDecomposition {
    using GraphOrderd = std::unordered_map<Node, std::deque<Node>>;

    const Node m_n;
    const std::vector<Node> m_size;
    const GraphOrderd m_tree;
    const std::vector<Node> m_height;
    const std::vector<std::pair<Node, Node>> m_root_par;
    const std::vector<Node> m_ids;
    const std::vector<Node> m_order;
    const std::vector<Node> m_edge_ids;

    static auto constructGraph(const Graph<Node, Cost>& tree) {
      auto n = tree.size();
      std::deque<std::pair<Node, Node>> order;
      std::vector<Node> used(n);
      std::stack<std::pair<Node, Node>> stk;
      stk.emplace(0, -1);
      used[0] = true;
      while (!stk.empty()) {
        auto [f, p] = stk.top();
        order.emplace_front(f, p);
        stk.pop();
        for (const auto& [t, _] : tree.getEdges(f)) {
          if (used[t]) {
            continue;
            ;
          }
          used[t] = true;
          stk.emplace(t, f);
        }
      }

      std::vector<Node> size(n, 1);
      GraphOrderd hld_tree;
      for (const auto& [f, p] : order) {
        Node size_sum = 1;
        Node size_max = 0;
        std::deque<Node> to_list;
        for (const auto& [t, _] : tree.getEdges(f)) {
          if (t == p) { continue; }
          if (size[t] > size_max) {
            size_max = size[t];
            to_list.emplace_back(t);
          } else {
            to_list.emplace_front(t);
          }
          size_sum += size[t];
        }
        if (!to_list.empty()) { hld_tree.emplace(f, to_list); }
        size[f] = size_sum;
      }
      return hld_tree;
    }

    static auto constructSize(const Graph<Node, Cost>& tree) {
      auto n = tree.size();
      std::deque<std::pair<Node, Node>> order;
      std::vector<Node> used(n);
      std::stack<std::pair<Node, Node>> stk;
      stk.emplace(0, -1);
      used[0] = true;
      while (!stk.empty()) {
        auto [f, p] = stk.top();
        order.emplace_front(f, p);
        stk.pop();
        for (const auto& [t, _] : tree.getEdges(f)) {
          if (used[t]) {
            continue;
            ;
          }
          used[t] = true;
          stk.emplace(t, f);
        }
      }

      std::vector<Node> size(n, 1);
      for (const auto& [f, p] : order) {
        Node size_sum = 1;
        for (const auto& [t, _] : tree.getEdges(f)) {
          if (t == p) { continue; }
          size_sum += size[t];
        }
        size[f] = size_sum;
      }
      return size;
    }

    static auto constructRootPar(Node n, const GraphOrderd& tree) {
      std::vector<std::pair<Node, Node>> root_par(n);
      std::stack<std::tuple<Node, Node, Node>> stk;
      stk.emplace(0, 0, -1);
      while (!stk.empty()) {
        auto [f, root, par] = stk.top();
        stk.pop();

        if (tree.find(f) == tree.end()) {
          root_par[f] = {root, par};
          continue;
        }
        auto itr = tree.at(f).rbegin();
        stk.emplace(*itr, root, par);
        root_par[f] = {root, par};
        for (++itr; itr != tree.at(f).rend(); ++itr) {
          stk.emplace(*itr, *itr, f);
        }
      }
      return root_par;
    }
    static auto constructHeight(Node n, const GraphOrderd& tree) {
      std::vector<Node> height(n);
      std::queue<Node> q;
      q.emplace(0);
      while (!q.empty()) {
        auto f = q.front();
        q.pop();
        if (tree.find(f) == tree.end()) { continue; }
        for (const auto& t : tree.at(f)) {
          height[t] = height[f] + 1;
          q.emplace(t);
        }
      }
      return height;
    }

    auto constructIds() const {
      std::vector<Node> ids(m_n);
      Node val = 0;
      std::stack<Node> stk;
      stk.emplace(0);
      while (!stk.empty()) {
        auto f = stk.top();
        stk.pop();
        ids[f] = val;
        ++val;
        if (m_tree.find(f) == m_tree.end()) { continue; }
        for (const auto& t : m_tree.at(f)) { stk.emplace(t); }
      }
      return ids;
    }

    auto constructOrder() const {
      std::vector<Node> order(m_n);
      for (int i = 0; i < m_n; ++i) { order[m_ids[i]] = i; }
      return order;
    }
    /*
     * 辺をnodeとして拡張した場合の辺nodeだけIDを振る
     * (1) - (2)
     * (1) - (e) - (2)
     * [-1, -1, 0]
     */
    auto constructEdgeIds() const {
      Node edge_size = (m_n >> 1);
      std::vector<Node> edge_ids(m_n, -1);
      Node val = 0;
      std::stack<Node> stk;
      stk.emplace(0);
      while (!stk.empty()) {
        auto f = stk.top();
        stk.pop();
        if (f > edge_size) {
          edge_ids[f] = val;
          ++val;
        }
        if (m_tree.find(f) == m_tree.end()) { continue; }
        for (const auto& t : m_tree.at(f)) { stk.emplace(t); }
      }
      return edge_ids;
    }

  public:
    HeavyLightDecomposition(const Graph<Node, Cost>& tree)
        : m_n(tree.size()),
          m_size(constructSize(tree)),
          m_tree(constructGraph(tree)),
          m_root_par(constructRootPar(m_n, m_tree)),
          m_height(constructHeight(m_n, m_tree)),
          m_ids(constructIds()),
          m_order(constructOrder()),
          m_edge_ids(constructEdgeIds()) {}

    auto getId(Node i) const { return m_ids[i]; }
    auto getEdgeId(Node i) const { return m_edge_ids[i]; }
    auto getOrder(Node i) const { return m_order[i]; }

    auto lca(Node f, Node t) const {
      do {
        auto [fr, fp] = m_root_par[f];
        auto [tr, tp] = m_root_par[t];
        if (fr == tr) { break; }
        auto fph = (fp > -1) ? m_height[fp] : -1;
        auto tph = (tp > -1) ? m_height[tp] : -1;
        if (fph < tph) {
          t = tp;
        } else {
          f = fp;
        }
      } while (true);
      return (m_height[f] < m_height[t]) ? f : t;
    }

    auto range(Node f, Node t) const {
      std::deque<std::pair<Node, Node>> ret;
      auto add = [&](Node from, Node to) {
        auto l = std::min(m_ids[from], m_ids[to]);
        auto r = std::max(m_ids[from], m_ids[to]);
        ret.emplace_back(l, r);
      };
      do {
        auto [fr, fp] = m_root_par[f];
        auto [tr, tp] = m_root_par[t];
        if (fr == tr) {
          add(f, t);
          break;
        }
        auto fph = (fp > -1) ? m_height[fp] : -1;
        auto tph = (tp > -1) ? m_height[tp] : -1;
        if (fph < tph) {
          add(t, tr);
          t = tp;
        } else {
          add(f, fr);
          f = fp;
        }
      } while (true);
      return ret;
    }

    auto rangeEdge(Node f, Node t) const {
      Node edge_size = (m_n >> 1);
      std::deque<std::pair<Node, Node>> ret;
      auto add = [&](Node from, Node to) {
        auto l = std::min(m_ids[from], m_ids[to]);
        auto r = std::max(m_ids[from], m_ids[to]);
        if (m_order[l] <= edge_size) { ++l; }
        if (m_order[r] <= edge_size) { --r; }
        if (l > r) { return; }
        auto edge_l = m_edge_ids[m_order[l]];
        auto edge_r = m_edge_ids[m_order[r]];
        ret.emplace_back(edge_l, edge_r);
      };
      do {
        auto [fr, fp] = m_root_par[f];
        auto [tr, tp] = m_root_par[t];
        if (fr == tr) {
          add(f, t);
          break;
        }
        auto fph = (fp > -1) ? m_height[fp] : -1;
        auto tph = (tp > -1) ? m_height[tp] : -1;
        if (fph < tph) {
          add(t, tr);
          t = tp;
        } else {
          add(f, fr);
          f = fp;
        }
      } while (true);
      return ret;
    }

    auto rangeSubTree(Node f) const {
      return std::pair<Node, Node>{m_ids[f], m_ids[f] + m_size[f] - 1};
    }
  };
}  // namespace mtd

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

namespace mtd {
  template <class Node, class Cost>
  class AuxiliaryTree {
    // 定数倍高速化のため破壊的

    std::vector<int> compres_map;

    const std::vector<Cost> depth_cost;
    const HeavyLightDecomposition<Node, Cost> hld;

    auto construct_depth(const Graph<Node, Cost>& tree) const {
      std::vector<Cost> _depth_cost(tree.size());
      std::vector<int> used(tree.size());
      auto dfs = [&](auto&& self, Node from) -> void {
        used[from] = true;
        for (const auto& [to, c] : tree.getEdges(from))
          if (!used[to]) {
            _depth_cost[to] = _depth_cost[from] + c;
            self(self, to);
          }
      };
      dfs(dfs, 0);
      return _depth_cost;
    }

  public:
    AuxiliaryTree(const Graph<Node, Cost>& tree)
        : compres_map(tree.size()),
          depth_cost(construct_depth(tree)),
          hld(tree) {}

    auto compression(const std::vector<Node>& nodes) {
      auto compare = [&](int a, int b) { return hld.getId(a) < hld.getId(b); };

      // 元の頂点集合

      auto nodes_set =
          std::set<int, decltype(compare)>(nodes.begin(), nodes.end(), compare);
      auto nodes_set_with_lca = nodes_set;

      // pre orderでの全ての隣接nodeのLCAを求める

      for (auto itr = nodes_set_with_lca.begin();
           std::next(itr) != nodes_set_with_lca.end(); ++itr) {
        nodes_set_with_lca.emplace(hld.lca(*itr, *std::next(itr)));
      }

      // 座標圧縮

      int at_size = nodes_set_with_lca.size();
      for (int i = 0; auto x : nodes_set_with_lca) {
        compres_map[x] = i;
        ++i;
      }

      // LCAを含めた全てのnodeで子孫関係を保って辺を張る

      std::stack<int> stk;
      Graph<Node, Cost> auxiliary_tree(at_size);
      for (auto nd : nodes_set_with_lca) {
        while (!stk.empty() && hld.lca(stk.top(), nd) != stk.top()) {
          stk.pop();
        }
        if (!stk.empty()) {
          auto f = compres_map[stk.top()];
          auto t = compres_map[nd];
          auto c = depth_cost[stk.top()] + depth_cost[nd] -
                   depth_cost[hld.lca(stk.top(), nd)] * 2;
          auxiliary_tree.addEdge(f, t, c);
        }
        stk.emplace(nd);
      }
      return auxiliary_tree;
    }
  };
}  // namespace mtd

Back to top page