Skip to content

lowest_common_ancestor.hpp

SECTIONGraph INCLUDEnoya/lowest_common_ancestor.hpp

Sparse-table-based LCA with O(n log n) build and O(1) query.

Verified by jump_on_tree, lca.

预处理有根树后回答两点最近公共祖先、距离和向上跳祖先;适合频繁树上定位。

Implementation

View on GitHub

#ifndef NOYA_LOWEST_COMMON_ANCESTOR_HPP
#define NOYA_LOWEST_COMMON_ANCESTOR_HPP 1

/// @complexity Time: O(n log n) build and O(1) LCA/ancestor queries.
/// Space: O(n log n).

#include "noya/sparse_table.hpp"
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <limits>
#include <utility>
#include <vector>

namespace noya {
/// @brief Sparse-table-based LCA with O(n log n) build and O(1) query.
struct fastlca {
  struct minimum_pair_semilattice {
    using value_type = std::pair<int, int>;
    static value_type unit() {
      return {std::numeric_limits<int>::max(), -1};
    }
    static value_type op(value_type a, value_type b) { return std::min(a, b); }
  };
  int n;
  std::vector<int> dfn;
  std::vector<int> d;
  std::vector<int64_t> len;
  std::vector<int> siz;
  sparse_table<minimum_pair_semilattice> ST;
  bool weighted;

  std::vector<int> fa;
  std::vector<std::vector<int>> ancestors;

  template <class T>
  fastlca(const std::vector<T> &g = {}, const bool &_weighted = false,
          const int &root = 0) {
    weighted = _weighted;
    if (!g.empty())
      build(g, root);
  };

  template <class T>
  void build(const std::vector<std::vector<T>> &g, const int &root = 0) {
    n = int(g.size());
    std::vector<std::vector<int>> g2(n);

    for (int u = 0; u < n; u++) {
      for (auto &[v, w] : g[u]) {
        g2[u].push_back(v);
      }
    }
    build(g2, root);

    len.assign(n, 0);
    auto dfs = [&](auto &self, int u) -> void {
      for (auto &[v, w] : g[u]) {
        if (fa[v] == u) {
          len[v] = len[u] + w;
          self(self, v);
        }
      }
    };
    dfs(dfs, root);
  }

  void build(const std::vector<std::vector<int>> &g = {}, const int &root = 0) {
    n = int(g.size());
    d.assign(n, 0);
    dfn.assign(n, -1);
    siz.assign(n, 0);
    fa.assign(n, -1);

    int idx = 0;
    std::vector<std::pair<int, int>> a;
    a.reserve(n);
    auto dfs = [&](auto &self, int u, int parent) -> void {
      fa[u] = parent;
      siz[u] = 1;
      dfn[u] = idx++;
      a.push_back(parent == -1 ? std::pair{-1, -1}
                               : std::pair{dfn[parent], parent});
      for (int v : g[u]) {
        if (v == parent) {
          continue;
        }
        d[v] = d[u] + 1;
        self(self, v, u);
        siz[u] += siz[v];
      }
    };
    dfs(dfs, root, -1);
    ST.build(a);
    int levels = 1;
    while ((1 << levels) <= std::max(1, n)) {
      levels++;
    }
    ancestors.assign(levels, std::vector<int>(n, -1));
    ancestors[0] = fa;
    for (int level = 1; level < levels; level++) {
      for (int vertex = 0; vertex < n; vertex++) {
        int middle = ancestors[level - 1][vertex];
        if (middle != -1) {
          ancestors[level][vertex] = ancestors[level - 1][middle];
        }
      }
    }
  }

  /// @brief Check if node a is in the subtree of node b.
  bool is_subtree(int a, int b) {
    if (dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b]) {
      return true;
    } else {
      return false;
    }
  }

  /// @brief Return the lowest common ancestor of nodes u and v.
  int lca(int u, int v) const {
    assert(0 <= u && u < n);
    assert(0 <= v && v < n);

    if (u == v) {
      return u;
    } else {
      int a = dfn[u];
      int b = dfn[v];
      if (a > b) {
        std::swap(a, b);
      }
      return ST.prod(a + 1, b + 1).second;
    }
  }

  /// @brief Return the LCA when the tree is re-rooted at c.
  int rooted_lca(int a, int b, int c) const {
    return lca(a, b) ^ lca(a, c) ^ lca(b, c);
  }

  /// @brief Return the k-th ancestor of vertex, or -1 when it does not exist.
  int kth_ancestor(int vertex, int k) const {
    assert(0 <= vertex && vertex < n);
    if (k < 0 || k > d[vertex]) {
      return -1;
    }
    for (int level = 0; k > 0; level++, k >>= 1) {
      if (k & 1) {
        vertex = ancestors[level][vertex];
      }
    }
    return vertex;
  }

  /// @brief Return the k-th vertex on path first->second, zero-indexed.
  int kth_on_path(int first, int second, int k) const {
    int ancestor = lca(first, second);
    int first_length = d[first] - d[ancestor];
    int second_length = d[second] - d[ancestor];
    if (k < 0 || k > first_length + second_length) {
      return -1;
    }
    if (k <= first_length) {
      return kth_ancestor(first, k);
    }
    return kth_ancestor(second, first_length + second_length - k);
  }

  /// @brief Return the (weighted or unweighted) distance between nodes a and b.
  int64_t distance(int a, int b) const {
    int c = lca(a, b);
    if (!weighted) {
      return d[a] + d[b] - d[c] * 2;
    } else {
      return len[a] + len[b] - len[c] * 2;
    }
  }

  /// @brief Compute the intersection of paths (a,b) and (c,d) as a pair of endpoints.
  std::pair<int, int> intersection(int a, int b, int c, int d) const {
    int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);
    int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);
    int x = ab ^ ac ^ bc;
    int y = ab ^ ad ^ bd;
    if (x != y) {
      return {x, y};
    }
    int z = ac ^ ad ^ cd;
    if (x != z) {
      x = -1;
    }
    return {x, x};
  }

  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);
  }

  /// @brief Return the list of vertices on the path from a to b.
  std::vector<int> path(int a, int b) {
    int c = lca(a, b);
    std::vector<int> ac;
    while (a != c) {
      ac.push_back(a);
      a = fa[a];
    }
    std::vector<int> bc;
    while (b != c) {
      bc.push_back(b);
      b = fa[b];
    }
    std::vector<int> res = std::move(ac);
    res.push_back(c);
    res.insert(res.end(), bc.rbegin(), bc.rend());
    return res;
  }
};
} // namespace noya

#endif // NOYA_LOWEST_COMMON_ANCESTOR_HPP
#include <algorithm>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

/// @complexity Time: O(n log n) build and O(1) LCA/ancestor queries.
/// Space: O(n log n).

/// @complexity Time: O(n log n) build and O(1) idempotent range query.
/// Space: O(n log n).

namespace noya {

namespace internal {

template <class Algebra, auto Operation, auto Identity,
          bool Legacy = !std::is_same_v<decltype(Operation), std::nullptr_t>>
struct sparse_table_algebra;

template <class Semilattice, auto Operation, auto Identity>
struct sparse_table_algebra<Semilattice, Operation, Identity, false> {
  static_assert(std::is_same_v<decltype(Identity), std::nullptr_t>,
                "a policy sparse_table gets its identity from unit()");

  using value_type = typename Semilattice::value_type;

  static value_type unit() { return Semilattice::unit(); }
  static value_type op(const value_type &left, const value_type &right) {
    return Semilattice::op(left, right);
  }
};

template <class Value, auto Operation, auto Identity>
struct sparse_table_algebra<Value, Operation, Identity, true> {
  using value_type = Value;

  static value_type unit() {
    if constexpr (std::is_same_v<decltype(Identity), std::nullptr_t>) {
      return value_type{};
    } else {
      return value_type(Identity);
    }
  }

  static value_type op(const value_type &left, const value_type &right) {
    return Operation(left, right);
  }
};

} // namespace internal

/// @brief Static O(1) range product for an idempotent semilattice.
/// The preferred interface is sparse_table<Semilattice>, where Semilattice
/// provides value_type, unit(), and an associative, idempotent op(). The legacy
/// sparse_table<Value, operation, identity> spelling remains supported.
/// @details For k = floor(log2(r-l)), [l, r) is covered by the length-2^k
/// blocks beginning at l and ending at r. They may overlap, so combining them
/// is correct precisely because op(x, x) = x.
template <class Algebra, auto Operation = nullptr, auto Identity = nullptr>
class sparse_table {
  using algebra_type =
      internal::sparse_table_algebra<Algebra, Operation, Identity>;

public:
  using value_type = typename algebra_type::value_type;

  sparse_table() = default;

  explicit sparse_table(const std::vector<value_type> &values) {
    build(values);
  }

  /// @brief Rebuild from a static array.
  void build(const std::vector<value_type> &values) {
    size_ = int(values.size());
    int levels = int(std::bit_width(static_cast<unsigned>(size_)));
    table_.assign(levels, {});
    if (size_ == 0) {
      return;
    }

    table_[0] = values;
    for (int level = 1; level < levels; level++) {
      int width = 1 << level;
      int half = width >> 1;
      table_[level].resize(size_ - width + 1);
      for (int left = 0; left + width <= size_; left++) {
        table_[level][left] = algebra_type::op(table_[level - 1][left],
                                               table_[level - 1][left + half]);
      }
    }
  }

  int size() const { return size_; }
  bool empty() const { return size_ == 0; }

  /// @brief Return the idempotent product over [left, right).
  value_type prod(int left, int right) const {
    assert(0 <= left && left <= right && right <= size_);
    if (left == right) {
      return algebra_type::unit();
    }
    int level = int(std::bit_width(static_cast<unsigned>(right - left))) - 1;
    int width = 1 << level;
    return algebra_type::op(table_[level][left], table_[level][right - width]);
  }

private:
  int size_ = 0;
  std::vector<std::vector<value_type>> table_;
};

} // namespace noya

namespace noya {
/// @brief Sparse-table-based LCA with O(n log n) build and O(1) query.
struct fastlca {
  struct minimum_pair_semilattice {
    using value_type = std::pair<int, int>;
    static value_type unit() {
      return {std::numeric_limits<int>::max(), -1};
    }
    static value_type op(value_type a, value_type b) { return std::min(a, b); }
  };
  int n;
  std::vector<int> dfn;
  std::vector<int> d;
  std::vector<int64_t> len;
  std::vector<int> siz;
  sparse_table<minimum_pair_semilattice> ST;
  bool weighted;

  std::vector<int> fa;
  std::vector<std::vector<int>> ancestors;

  template <class T>
  fastlca(const std::vector<T> &g = {}, const bool &_weighted = false,
          const int &root = 0) {
    weighted = _weighted;
    if (!g.empty())
      build(g, root);
  };

  template <class T>
  void build(const std::vector<std::vector<T>> &g, const int &root = 0) {
    n = int(g.size());
    std::vector<std::vector<int>> g2(n);

    for (int u = 0; u < n; u++) {
      for (auto &[v, w] : g[u]) {
        g2[u].push_back(v);
      }
    }
    build(g2, root);

    len.assign(n, 0);
    auto dfs = [&](auto &self, int u) -> void {
      for (auto &[v, w] : g[u]) {
        if (fa[v] == u) {
          len[v] = len[u] + w;
          self(self, v);
        }
      }
    };
    dfs(dfs, root);
  }

  void build(const std::vector<std::vector<int>> &g = {}, const int &root = 0) {
    n = int(g.size());
    d.assign(n, 0);
    dfn.assign(n, -1);
    siz.assign(n, 0);
    fa.assign(n, -1);

    int idx = 0;
    std::vector<std::pair<int, int>> a;
    a.reserve(n);
    auto dfs = [&](auto &self, int u, int parent) -> void {
      fa[u] = parent;
      siz[u] = 1;
      dfn[u] = idx++;
      a.push_back(parent == -1 ? std::pair{-1, -1}
                               : std::pair{dfn[parent], parent});
      for (int v : g[u]) {
        if (v == parent) {
          continue;
        }
        d[v] = d[u] + 1;
        self(self, v, u);
        siz[u] += siz[v];
      }
    };
    dfs(dfs, root, -1);
    ST.build(a);
    int levels = 1;
    while ((1 << levels) <= std::max(1, n)) {
      levels++;
    }
    ancestors.assign(levels, std::vector<int>(n, -1));
    ancestors[0] = fa;
    for (int level = 1; level < levels; level++) {
      for (int vertex = 0; vertex < n; vertex++) {
        int middle = ancestors[level - 1][vertex];
        if (middle != -1) {
          ancestors[level][vertex] = ancestors[level - 1][middle];
        }
      }
    }
  }

  /// @brief Check if node a is in the subtree of node b.
  bool is_subtree(int a, int b) {
    if (dfn[b] <= dfn[a] && dfn[a] < dfn[b] + siz[b]) {
      return true;
    } else {
      return false;
    }
  }

  /// @brief Return the lowest common ancestor of nodes u and v.
  int lca(int u, int v) const {
    assert(0 <= u && u < n);
    assert(0 <= v && v < n);

    if (u == v) {
      return u;
    } else {
      int a = dfn[u];
      int b = dfn[v];
      if (a > b) {
        std::swap(a, b);
      }
      return ST.prod(a + 1, b + 1).second;
    }
  }

  /// @brief Return the LCA when the tree is re-rooted at c.
  int rooted_lca(int a, int b, int c) const {
    return lca(a, b) ^ lca(a, c) ^ lca(b, c);
  }

  /// @brief Return the k-th ancestor of vertex, or -1 when it does not exist.
  int kth_ancestor(int vertex, int k) const {
    assert(0 <= vertex && vertex < n);
    if (k < 0 || k > d[vertex]) {
      return -1;
    }
    for (int level = 0; k > 0; level++, k >>= 1) {
      if (k & 1) {
        vertex = ancestors[level][vertex];
      }
    }
    return vertex;
  }

  /// @brief Return the k-th vertex on path first->second, zero-indexed.
  int kth_on_path(int first, int second, int k) const {
    int ancestor = lca(first, second);
    int first_length = d[first] - d[ancestor];
    int second_length = d[second] - d[ancestor];
    if (k < 0 || k > first_length + second_length) {
      return -1;
    }
    if (k <= first_length) {
      return kth_ancestor(first, k);
    }
    return kth_ancestor(second, first_length + second_length - k);
  }

  /// @brief Return the (weighted or unweighted) distance between nodes a and b.
  int64_t distance(int a, int b) const {
    int c = lca(a, b);
    if (!weighted) {
      return d[a] + d[b] - d[c] * 2;
    } else {
      return len[a] + len[b] - len[c] * 2;
    }
  }

  /// @brief Compute the intersection of paths (a,b) and (c,d) as a pair of endpoints.
  std::pair<int, int> intersection(int a, int b, int c, int d) const {
    int ab = lca(a, b), ac = lca(a, c), ad = lca(a, d);
    int bc = lca(b, c), bd = lca(b, d), cd = lca(c, d);
    int x = ab ^ ac ^ bc;
    int y = ab ^ ad ^ bd;
    if (x != y) {
      return {x, y};
    }
    int z = ac ^ ad ^ cd;
    if (x != z) {
      x = -1;
    }
    return {x, x};
  }

  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);
  }

  /// @brief Return the list of vertices on the path from a to b.
  std::vector<int> path(int a, int b) {
    int c = lca(a, b);
    std::vector<int> ac;
    while (a != c) {
      ac.push_back(a);
      a = fa[a];
    }
    std::vector<int> bc;
    while (b != c) {
      bc.push_back(b);
      b = fa[b];
    }
    std::vector<int> res = std::move(ac);
    res.push_back(c);
    res.insert(res.end(), bc.rbegin(), bc.rend());
    return res;
  }
};
} // namespace noya