tests/lca.test.cpp

View this file on GitHub
Last update: 2025-02-10 20:05:30+09:00
Problem: https://judge.yosupo.jp/problem/lca

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/lca

#include <iostream>
#include "tools/lca.hpp"

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, Q;
  std::cin >> N >> Q;
  tools::lca lca(N);
  for (int i = 1; i < N; ++i) {
    int p;
    std::cin >> p;
    lca.add_edge(i, p);
  }
  lca.build(0);

  for (int q = 0; q < Q; ++q) {
    int u, v;
    std::cin >> u >> v;
    std::cout << lca.query(u, v) << '\n';
  }

  return 0;
}

#line 1 "tests/lca.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/lca

#include <iostream>
#line 1 "tools/lca.hpp"



#include <cstdint>
#include <vector>
#include <cstddef>
#include <cassert>
#include <numeric>
#include <limits>
#include <stack>
#include <utility>
#include <algorithm>
#include <tuple>
#line 1 "tools/ceil.hpp"



#line 5 "tools/ceil.hpp"
#include <type_traits>
#line 1 "tools/is_integral.hpp"



#line 5 "tools/is_integral.hpp"

namespace tools {
  template <typename T>
  struct is_integral : ::std::is_integral<T> {};

  template <typename T>
  inline constexpr bool is_integral_v = ::tools::is_integral<T>::value;
}


#line 1 "tools/is_unsigned.hpp"



#line 5 "tools/is_unsigned.hpp"

namespace tools {
  template <typename T>
  struct is_unsigned : ::std::is_unsigned<T> {};

  template <typename T>
  inline constexpr bool is_unsigned_v = ::tools::is_unsigned<T>::value;
}


#line 8 "tools/ceil.hpp"

namespace tools {
  template <typename M, typename N> requires (
    ::tools::is_integral_v<M> && !::std::is_same_v<::std::remove_cv_t<M>, bool> &&
    ::tools::is_integral_v<N> && !::std::is_same_v<::std::remove_cv_t<N>, bool>)
  constexpr ::std::common_type_t<M, N> ceil(const M x, const N y) noexcept {
    assert(y != 0);
    if (y >= 0) {
      if (x > 0) {
        return (x - 1) / y + 1;
      } else {
        if constexpr (::tools::is_unsigned_v<::std::common_type_t<M, N>>) {
          return 0;
        } else {
          return x / y;
        }
      }
    } else {
      if (x >= 0) {
        if constexpr (::tools::is_unsigned_v<::std::common_type_t<M, N>>) {
          return 0;
        } else {
          return x / y;
        }
      } else {
        return (x + 1) / y + 1;
      }
    }
  }
}


#line 1 "tools/less_by.hpp"



namespace tools {

  template <class F>
  class less_by {
  private:
    F selector;

  public:
    less_by(const F& selector) : selector(selector) {
    }

    template <class T>
    bool operator()(const T& x, const T& y) const {
      return selector(x) < selector(y);
    }
  };
}


#line 1 "tools/ceil_log2.hpp"



#line 1 "tools/bit_width.hpp"



#include <bit>
#line 1 "tools/is_signed.hpp"



#line 5 "tools/is_signed.hpp"

namespace tools {
  template <typename T>
  struct is_signed : ::std::is_signed<T> {};

  template <typename T>
  inline constexpr bool is_signed_v = ::tools::is_signed<T>::value;
}


#line 1 "tools/make_unsigned.hpp"



#line 5 "tools/make_unsigned.hpp"

namespace tools {
  template <typename T>
  struct make_unsigned : ::std::make_unsigned<T> {};

  template <typename T>
  using make_unsigned_t = typename ::tools::make_unsigned<T>::type;
}


#line 10 "tools/bit_width.hpp"

namespace tools {
  template <typename T>
  constexpr int bit_width(T) noexcept;

  template <typename T>
  constexpr int bit_width(const T x) noexcept {
    static_assert(::tools::is_integral_v<T> && !::std::is_same_v<::std::remove_cv_t<T>, bool>);
    if constexpr (::tools::is_signed_v<T>) {
      assert(x >= 0);
      return ::tools::bit_width<::tools::make_unsigned_t<T>>(x);
    } else {
      return ::std::bit_width(x);
    }
  }
}


#line 6 "tools/ceil_log2.hpp"

namespace tools {
  template <typename T>
  constexpr T ceil_log2(T x) noexcept {
    assert(x > 0);
    return ::tools::bit_width(x - 1);
  }
}


#line 1 "tools/floor_log2.hpp"



#line 6 "tools/floor_log2.hpp"

namespace tools {
  template <typename T>
  constexpr T floor_log2(T x) noexcept {
    assert(x > 0);
    return ::tools::bit_width(x) - 1;
  }
}


#line 1 "tools/pow2.hpp"



#line 6 "tools/pow2.hpp"

namespace tools {

  template <typename T, typename ::std::enable_if<::std::is_unsigned<T>::value, ::std::nullptr_t>::type = nullptr>
  constexpr T pow2(const T x) {
    return static_cast<T>(1) << x;
  }

  template <typename T, typename ::std::enable_if<::std::is_signed<T>::value, ::std::nullptr_t>::type = nullptr>
  constexpr T pow2(const T x) {
    return static_cast<T>(static_cast<typename ::std::make_unsigned<T>::type>(1) << static_cast<typename ::std::make_unsigned<T>::type>(x));
  }
}


#line 19 "tools/lca.hpp"

namespace tools {
  class lca {
    using u32 = ::std::uint32_t;
    ::std::vector<::std::vector<u32>> m_graph;
    ::std::vector<u32> m_depth;
    ::std::vector<u32> m_tour;
    ::std::vector<u32> m_in;
    u32 m_block_size;
    ::std::vector<::std::vector<u32>> m_sparse_table;
    ::std::vector<::std::vector<::std::vector<u32>>> m_lookup_table;
    ::std::vector<u32> m_patterns;

    bool built() const {
      return !this->m_depth.empty();
    }

    u32 nblocks() const {
      return ::tools::ceil(this->m_tour.size(), this->m_block_size);
    }

    auto less_by_depth() const {
      return ::tools::less_by([&](const auto v) { return this->m_depth[v]; });
    }

  public:
    lca() = default;
    explicit lca(const ::std::size_t n) : m_graph(n) {
      assert(n >= 1);
    }

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

    void add_edge(const ::std::size_t u, const ::std::size_t v) {
      assert(!this->built());
      assert(u < this->size());
      assert(v < this->size());
      assert(u != v);
      this->m_graph[u].push_back(v);
      this->m_graph[v].push_back(u);
    }

    void build(const ::std::size_t r) {
      assert(!this->built());
      assert(::std::accumulate(this->m_graph.begin(), this->m_graph.end(), static_cast<::std::size_t>(0), [](const auto sum, const auto& neighbors) { return sum + neighbors.size(); }) == 2 * (this->size() - 1));

      this->m_depth.assign(this->size(), ::std::numeric_limits<u32>::max());
      this->m_tour.resize(2 * this->size() - 1);
      this->m_in.resize(this->size());

      u32 t = 0;
      ::std::stack<::std::pair<u32, u32>> stack;
      stack.emplace(r, 0);
      while (!stack.empty()) {
        const auto [here, depth] = stack.top();
        stack.pop();
        this->m_tour[t] = here;
        if (this->m_depth[here] == ::std::numeric_limits<u32>::max()) {
          this->m_depth[here] = depth;
          this->m_in[here] = t;
          for (const auto next : this->m_graph[here]) {
            if (this->m_depth[next] == ::std::numeric_limits<u32>::max()) {
              stack.emplace(here, depth);
              stack.emplace(next, depth + 1);
            }
          }
        }
        ++t;
      }

      if (this->size() > 1) {
        this->m_tour.pop_back();
      }

      this->m_block_size = ::std::max<u32>(1, ::tools::ceil(::tools::ceil_log2(this->m_tour.size()), 2));
      this->m_sparse_table.resize(::tools::floor_log2(this->nblocks()) + 1);
      this->m_sparse_table[0].resize(this->nblocks());
      for (u32 b = 0; (b + 1) * this->m_block_size <= this->m_tour.size(); ++b) {
        const auto l = b * this->m_block_size;
        const auto r = ::std::min<u32>(l + this->m_block_size, this->m_tour.size());
        this->m_sparse_table[0][b] = *::std::min_element(this->m_tour.begin() + l, this->m_tour.begin() + r, this->less_by_depth());
      }
      for (u32 h = 1; h < this->m_sparse_table.size(); ++h) {
        this->m_sparse_table[h].resize(this->nblocks() + UINT32_C(1) - (UINT32_C(1) << h));
        for (u32 b = 0; b < this->m_sparse_table[h].size(); ++b) {
          this->m_sparse_table[h][b] = ::std::min(this->m_sparse_table[h - 1][b], this->m_sparse_table[h - 1][b + (UINT32_C(1) << (h - 1))], this->less_by_depth());
        }
      }

      this->m_lookup_table.resize(::tools::pow2(this->m_block_size - 1));
      for (u32 p = 0; p < this->m_lookup_table.size(); ++p) {
        this->m_lookup_table[p].resize(this->m_block_size + 1);
        for (u32 l = 0; l <= this->m_block_size; ++l) {
          this->m_lookup_table[p][l].resize(this->m_block_size + 1);
        }

        ::std::vector<u32> partial_sum(this->m_block_size);
        partial_sum[0] = this->m_block_size;
        for (u32 i = 1; i < this->m_block_size; ++i) {
          partial_sum[i] = partial_sum[i - 1] - UINT32_C(1) + (((p >> (i - 1)) & UINT32_C(1)) << 1);
        }

        for (u32 l = 0; l < this->m_block_size; ++l) {
          this->m_lookup_table[p][l][l + 1] = l;
          for (u32 r = l + 2; r <= this->m_block_size; ++r) {
            this->m_lookup_table[p][l][r] = ::std::min(this->m_lookup_table[p][l][r - 1], r - 1, ::tools::less_by([&](const auto i) { return partial_sum[i]; }));
          }
        }
      }

      this->m_patterns.resize(this->nblocks());
      for (u32 b = 0; b * this->m_block_size < this->m_tour.size(); ++b) {
        const auto l = b * this->m_block_size;
        const auto r = ::std::min<u32>(l + this->m_block_size, this->m_tour.size());
        this->m_patterns[b] = 0;
        for (u32 i = l; i + 1 < r; ++i) {
          this->m_patterns[b] |= static_cast<u32>(this->m_depth[this->m_tour[i]] < this->m_depth[this->m_tour[i + 1]]) << (i - l);
        }
      }
    }

    ::std::size_t depth(const ::std::size_t v) const {
      assert(this->built());
      assert(v < this->size());
      return this->m_depth[v];
    }

    ::std::size_t query(::std::size_t u, ::std::size_t v) const {
      assert(this->built());
      assert(u < this->size());
      assert(v < this->size());

      ::std::tie(u, v) = ::std::minmax({u, v}, ::tools::less_by([&](const auto w) { return this->m_in[w]; }));

      const auto l = this->m_in[u];
      const auto r = this->m_in[v] + UINT32_C(1);
      const auto bl = ::tools::ceil(l, this->m_block_size);
      const auto br = r / this->m_block_size;
      u32 lca;
      if (br < bl) {
        lca = this->m_tour[br * this->m_block_size + this->m_lookup_table[this->m_patterns[br]][l % this->m_block_size][r % this->m_block_size]];
      } else {
        lca = u;
        if (bl < br) {
          const auto h = ::tools::floor_log2(br - bl);
          lca = ::std::min(this->m_sparse_table[h][bl], this->m_sparse_table[h][br - (UINT32_C(1) << h)], this->less_by_depth());
        }
        if (l < bl * this->m_block_size) {
          lca = ::std::min(lca, this->m_tour[(bl - UINT32_C(1)) * this->m_block_size + this->m_lookup_table[this->m_patterns[bl - 1]][l % this->m_block_size][this->m_block_size]], this->less_by_depth());
        }
        if (br * this->m_block_size < r) {
          lca = ::std::min(lca, this->m_tour[br * this->m_block_size + this->m_lookup_table[this->m_patterns[br]][0][r % this->m_block_size]], this->less_by_depth());
        }
      }

      return lca;
    }

    // for tools::auxiliary_tree
    ::std::size_t internal_in(const ::std::size_t v) const {
      assert(this->built());
      assert(v < this->size());
      return this->m_in[v];
    }
  };
}


#line 5 "tests/lca.test.cpp"

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, Q;
  std::cin >> N >> Q;
  tools::lca lca(N);
  for (int i = 1; i < N; ++i) {
    int p;
    std::cin >> p;
    lca.add_edge(i, p);
  }
  lca.build(0);

  for (int q = 0; q < Q; ++q) {
    int u, v;
    std::cin >> u >> v;
    std::cout << lca.query(u, v) << '\n';
  }

  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	almost_line_00	AC	220 ms	46 MB
g++	almost_line_01	AC	233 ms	46 MB
g++	binary_00	AC	287 ms	46 MB
g++	binary_01	AC	288 ms	46 MB
g++	binary_02	AC	270 ms	46 MB
g++	example_00	AC	5 ms	3 MB
g++	line_00	AC	151 ms	36 MB
g++	line_01	AC	165 ms	43 MB
g++	line_02	AC	85 ms	8 MB
g++	line_03	AC	76 ms	40 MB
g++	line_04	AC	67 ms	27 MB
g++	max_line_00	AC	194 ms	46 MB
g++	max_line_01	AC	192 ms	46 MB
g++	max_line_02	AC	194 ms	46 MB
g++	max_random_00	AC	334 ms	46 MB
g++	max_random_01	AC	326 ms	46 MB
g++	max_random_02	AC	353 ms	46 MB
g++	path_graph_root_centroid_00	AC	192 ms	46 MB
g++	path_graph_root_centroid_01	AC	186 ms	46 MB
g++	path_graph_root_centroid_02	AC	189 ms	46 MB
g++	random_00	AC	263 ms	37 MB
g++	random_01	AC	313 ms	43 MB
g++	random_02	AC	97 ms	8 MB
g++	random_03	AC	173 ms	40 MB
g++	random_04	AC	109 ms	27 MB