Skip to content

virtual_tree.hpp

SECTIONGraph INCLUDEnoya/graph/virtual_tree.hpp

把一次询问涉及的若干关键点及其 LCA 压成小树;适合多组树上关键点 DP。

Complexity: Time: O(k log k + k L) for k key vertices and LCA query cost L. Space: O(k).

跳到代码 · GitHub ↗

Implementation

当前头文件,省略 include guard;依赖见 #include。

/// @complexity Time: O(k log k + k L) for k key vertices and LCA query cost L.
/// Space: O(k).

#include "noya/graph/tree_preprocessing.hpp"

#include <algorithm>
#include <cassert>
#include <vector>

namespace noya {

/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
  std::vector<int> vs;
  std::vector<int> fa;
  std::vector<int> len;
  std::vector<bool> key;
};

/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
/// Accepts fastlca (L=O(1)) or hld (L=O(log n)) without copying the tree.
template <class Tree>
virtual_tree_result build_virtual_tree(const Tree &tr, std::vector<int> vs1) {
  virtual_tree_result res;
  if (vs1.empty()) {
    return res;
  }
  for (int u : vs1) {
    assert(0 <= u && u < tr.n);
  }
  std::sort(vs1.begin(), vs1.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
  res.vs = vs1;
  for (int i = 1; i < int(vs1.size()); i++) {
    res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
  }
  std::sort(res.vs.begin(), res.vs.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());

  int sz = int(res.vs.size());
  res.fa.assign(sz, -1);
  res.len.assign(sz, 0);
  res.key.resize(sz);
  std::vector<int> stk;
  int next_key = 0;
  for (int idx = 0; idx < sz; idx++) {
    int u = res.vs[idx];
    if (next_key < int(vs1.size()) && vs1[next_key] == u) {
      res.key[idx] = true;
      next_key++;
    }
    while (!stk.empty() && !tr.is_subtree(u, res.vs[stk.back()])) {
      stk.pop_back();
    }
    if (!stk.empty()) {
      res.fa[idx] = stk.back();
      res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
    }
    stk.push_back(idx);
  }
  return res;
}

} // namespace noya

 // NOYA_VIRTUAL_TREE_HPP
#ifndef NOYA_VIRTUAL_TREE_HPP
#define NOYA_VIRTUAL_TREE_HPP 1

/// @complexity Time: O(k log k + k L) for k key vertices and LCA query cost L.
/// Space: O(k).

#include "noya/graph/tree_preprocessing.hpp"

#include <algorithm>
#include <cassert>
#include <vector>

namespace noya {

/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
  std::vector<int> vs;
  std::vector<int> fa;
  std::vector<int> len;
  std::vector<bool> key;
};

/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
/// Accepts fastlca (L=O(1)) or hld (L=O(log n)) without copying the tree.
template <class Tree>
virtual_tree_result build_virtual_tree(const Tree &tr, std::vector<int> vs1) {
  virtual_tree_result res;
  if (vs1.empty()) {
    return res;
  }
  for (int u : vs1) {
    assert(0 <= u && u < tr.n);
  }
  std::sort(vs1.begin(), vs1.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
  res.vs = vs1;
  for (int i = 1; i < int(vs1.size()); i++) {
    res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
  }
  std::sort(res.vs.begin(), res.vs.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());

  int sz = int(res.vs.size());
  res.fa.assign(sz, -1);
  res.len.assign(sz, 0);
  res.key.resize(sz);
  std::vector<int> stk;
  int next_key = 0;
  for (int idx = 0; idx < sz; idx++) {
    int u = res.vs[idx];
    if (next_key < int(vs1.size()) && vs1[next_key] == u) {
      res.key[idx] = true;
      next_key++;
    }
    while (!stk.empty() && !tr.is_subtree(u, res.vs[stk.back()])) {
      stk.pop_back();
    }
    if (!stk.empty()) {
      res.fa[idx] = stk.back();
      res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
    }
    stk.push_back(idx);
  }
  return res;
}

} // namespace noya

#endif // NOYA_VIRTUAL_TREE_HPP
/// @complexity Time: O(k log k + k L) for k key vertices and LCA query cost L.
/// Space: O(k).

/// @complexity Time: O(n) build; common queries add O(1) work to their LCA calls,
/// except path(), which takes O(path length). Space: O(n).

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

namespace noya {

/// @brief Iterative rooted-tree preprocessing shared by RMQ LCA and HLD.
/// The graph is a connected tree, with integer endpoints or (endpoint, weight)
/// pairs. Root depth is zero, root parent is -1, and subtree ranges are [l, r).
/// No adjacency-list copy, RMQ table, jump table, or heavy paths are retained.
struct tree_preprocessing {
  int n = 0;
  int rt = -1;
  std::vector<int> fa, d, dfn, siz, tl;

  tree_preprocessing() = default;
  template <class Edge>
  explicit tree_preprocessing(const std::vector<std::vector<Edge>> &g,
                              int root = 0) {
    build(g, root);
  }

  template <class Edge> static int endpoint(const Edge &edge) {
    if constexpr (std::is_integral_v<Edge>) {
      return int(edge);
    } else {
      return int(std::get<0>(edge));
    }
  }

  template <class Edge>
  void build(const std::vector<std::vector<Edge>> &g, int root = 0) {
    n = int(g.size());
    rt = n == 0 ? -1 : root;
    fa.assign(n, -1);
    d.assign(n, 0);
    dfn.assign(n, -1);
    siz.assign(n, 1);
    tl.clear();
    tl.reserve(n);
    if (n == 0) {
      return;
    }
    assert(0 <= root && root < n);
    std::vector<int> stack{root};
    while (!stack.empty()) {
      int u = stack.back();
      stack.pop_back();
      assert(dfn[u] == -1);
      dfn[u] = int(tl.size());
      tl.push_back(u);
      for (auto it = g[u].rbegin(); it != g[u].rend(); ++it) {
        int v = endpoint(*it);
        assert(0 <= v && v < n);
        if (v != fa[u]) {
          fa[v] = u;
          d[v] = d[u] + 1;
          stack.push_back(v);
        }
      }
    }
    assert(int(tl.size()) == n);
    for (auto it = tl.rbegin(); it != tl.rend(); ++it) {
      if (fa[*it] != -1) {
        siz[fa[*it]] += siz[*it];
      }
    }
  }

  void build(const std::vector<std::vector<int>> &g = {}, int root = 0) {
    build<int>(g, root);
  }

  bool is_subtree(int a, int b) const {
    assert(0 <= a && a < n && 0 <= b && b < n);
    return dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b];
  }

  std::array<int, 2> subtree(int u) const {
    assert(0 <= u && u < n);
    return {dfn[u], dfn[u] + siz[u]};
  }
};

/// @brief Shared tree queries with compile-time dispatch to Derived::lca().
template <class Derived> struct tree_queries : tree_preprocessing {
  int rooted_lca(int a, int b, int root) const {
    const auto &self = static_cast<const Derived &>(*this);
    return self.lca(a, b) ^ self.lca(a, root) ^ self.lca(b, root);
  }

  std::int64_t distance(int a, int b) const {
    const auto &self = static_cast<const Derived &>(*this);
    int c = self.lca(a, b);
    return std::int64_t(d[a]) + d[b] - 2 * std::int64_t(d[c]);
  }

  std::pair<int, int> intersection(int a, int b, int c, int e) const {
    const auto &self = static_cast<const Derived &>(*this);
    int ab = self.lca(a, b), ac = self.lca(a, c), ae = self.lca(a, e);
    int bc = self.lca(b, c), be = self.lca(b, e), ce = self.lca(c, e);
    int x = ab ^ ac ^ bc;
    int y = ab ^ ae ^ be;
    if (x != y) {
      return {x, y};
    }
    int z = ac ^ ae ^ ce;
    return x == z ? std::pair{x, x} : std::pair{-1, -1};
  }

  std::pair<int, int> intersection(std::pair<int, int> a,
                                  std::pair<int, int> b) const {
    return intersection(a.first, a.second, b.first, b.second);
  }

  std::vector<int> path(int a, int b) const {
    const auto &self = static_cast<const Derived &>(*this);
    int c = self.lca(a, b);
    std::vector<int> left, right;
    while (a != c) {
      left.push_back(a);
      a = fa[a];
    }
    while (b != c) {
      right.push_back(b);
      b = fa[b];
    }
    left.push_back(c);
    left.insert(left.end(), right.rbegin(), right.rend());
    return left;
  }
};

} // namespace noya

#include <algorithm>
#include <cassert>
#include <vector>

namespace noya {

/// @brief Vertices and parent edges of a compressed virtual tree.
struct virtual_tree_result {
  std::vector<int> vs;
  std::vector<int> fa;
  std::vector<int> len;
  std::vector<bool> key;
};

/// @brief Build the minimal rooted tree containing all key vertices and LCAs.
/// Accepts fastlca (L=O(1)) or hld (L=O(log n)) without copying the tree.
template <class Tree>
virtual_tree_result build_virtual_tree(const Tree &tr, std::vector<int> vs1) {
  virtual_tree_result res;
  if (vs1.empty()) {
    return res;
  }
  for (int u : vs1) {
    assert(0 <= u && u < tr.n);
  }
  std::sort(vs1.begin(), vs1.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  vs1.erase(std::unique(vs1.begin(), vs1.end()), vs1.end());
  res.vs = vs1;
  for (int i = 1; i < int(vs1.size()); i++) {
    res.vs.push_back(tr.lca(vs1[i - 1], vs1[i]));
  }
  std::sort(res.vs.begin(), res.vs.end(),
            [&](int a, int b) { return tr.dfn[a] < tr.dfn[b]; });
  res.vs.erase(std::unique(res.vs.begin(), res.vs.end()), res.vs.end());

  int sz = int(res.vs.size());
  res.fa.assign(sz, -1);
  res.len.assign(sz, 0);
  res.key.resize(sz);
  std::vector<int> stk;
  int next_key = 0;
  for (int idx = 0; idx < sz; idx++) {
    int u = res.vs[idx];
    if (next_key < int(vs1.size()) && vs1[next_key] == u) {
      res.key[idx] = true;
      next_key++;
    }
    while (!stk.empty() && !tr.is_subtree(u, res.vs[stk.back()])) {
      stk.pop_back();
    }
    if (!stk.empty()) {
      res.fa[idx] = stk.back();
      res.len[idx] = tr.d[u] - tr.d[res.vs[stk.back()]];
    }
    stk.push_back(idx);
  }
  return res;
}

} // namespace noya

 // NOYA_VIRTUAL_TREE_HPP