This documentation is automatically generated by online-judge-tools/verification-helper
View the Project on GitHub suisen-cp/cp-library-cpp
#define PROBLEM "https://judge.yosupo.jp/problem/jump_on_tree" #include <iostream> #include "library/tree/lowest_common_ancestor.hpp" #include "library/tree/level_ancestor.hpp" int main() { std::ios::sync_with_stdio(false); std::cin.tie(nullptr); int n, q; std::cin >> n >> q; std::vector<std::vector<int>> g(n); for (int i = 0; i < n - 1; ++i) { int u, v; std::cin >> u >> v; g[u].push_back(v); g[v].push_back(u); } suisen::LowestCommonAncestor lca(g); suisen::LevelAncestor la(g); while (q --> 0) { int u, v, d; std::cin >> u >> v >> d; const int a = lca(u, v); const int du = la.get_depth(u) - la.get_depth(a); const int dv = la.get_depth(v) - la.get_depth(a); if (d <= du) { std::cout << la(u, d) << '\n'; } else if (d <= du + dv) { std::cout << la(v, du + dv - d) << '\n'; } else { std::cout << -1 << '\n'; } } return 0; }
#line 1 "test/src/tree/level_ancestor/jump_on_tree.test.cpp" #define PROBLEM "https://judge.yosupo.jp/problem/jump_on_tree" #include <iostream> #line 1 "library/tree/lowest_common_ancestor.hpp" #line 1 "library/algorithm/rmq_pm1_with_index.hpp" #include <cassert> #include <cstdint> #include <cmath> #include <algorithm> #line 1 "library/datastructure/sparse_table.hpp" #include <vector> namespace suisen { template <typename T, T(*op)(T, T), T(*e)()> struct SparseTable { SparseTable() = default; SparseTable(std::vector<T>&& a) : n(a.size()), log(floor_log2(n)), table(log + 1), flog(n + 1, 0) { build_table(std::move(a)); build_flog_table(); } SparseTable(const std::vector<T>& a) : SparseTable(std::vector<T>(a)) {} T operator()(int l, int r) const { if (l >= r) return e(); int i = flog[r - l]; return op(table[i][l], table[i][r - (1 << i)]); } T prod(int l, int r) const { return (*this)(l, r); } private: int n; int log; std::vector<std::vector<T>> table; std::vector<int> flog; void build_table(std::vector<T>&& a) { table[0] = std::move(a); for (int i = 0; i < log; ++i) { int lmax = n - (1 << (i + 1)); table[i + 1].resize(lmax + 1); for (int l = 0; l <= lmax; ++l) table[i + 1][l] = op(table[i][l], table[i][l + (1 << i)]); } } void build_flog_table() { for (int l = 0; l < log; ++l) { std::fill(flog.begin() + (1 << l), flog.begin() + (1 << (l + 1)), l); } std::fill(flog.begin() + (1 << log), flog.end(), log); } static int floor_log2(int i) { return 31 - __builtin_clz(i); } }; } // namespace suisen #line 10 "library/algorithm/rmq_pm1_with_index.hpp" namespace suisen { template <bool is_min_query = true> class RMQpm1WithIndex { static constexpr typename std::conditional_t<is_min_query, std::less<std::pair<int, int>>, std::greater<std::pair<int, int>>> comp{}; static constexpr typename std::conditional_t<is_min_query, std::less<int>, std::greater<int>> comp_val{}; static constexpr std::pair<int, int> op(std::pair<int, int> x, std::pair<int, int> y) { return comp(x, y) ? x : y; } static constexpr std::pair<int, int> e() { if constexpr (is_min_query) { return { std::numeric_limits<int>::max(), -1 }; } else { return { std::numeric_limits<int>::min(), -1 }; } } static constexpr int LOG = 4; static constexpr int SIZE = 1 << LOG; static constexpr class S { public: int prd[1 << RMQpm1WithIndex<is_min_query>::SIZE]; int arg[1 << RMQpm1WithIndex<is_min_query>::SIZE]; constexpr S() : prd(), arg(), sum() { prd[is_min_query] = sum[is_min_query] = -1, prd[not is_min_query] = sum[not is_min_query] = 1; arg[is_min_query] = arg[not is_min_query] = 0; for (int n = 2; n <= RMQpm1WithIndex<is_min_query>::SIZE; n <<= 1) { for (int s = (1 << n) - 1; s >= 0; --s) merge(s, n >> 1); } } private: int sum[1 << RMQpm1WithIndex<is_min_query>::SIZE]; constexpr void merge(int s, int half) { int upper = s >> half, lower = s ^ (upper << half); if (RMQpm1WithIndex<is_min_query>::comp_val(prd[lower], sum[lower] + prd[upper])) { prd[s] = prd[lower]; arg[s] = arg[lower]; } else { prd[s] = sum[lower] + prd[upper]; arg[s] = half + arg[upper]; } sum[s] = sum[lower] + sum[upper]; } } tabs{}; public: RMQpm1WithIndex(std::vector<int>&& x) : n(x.size()), m((n + SIZE - 1) >> LOG), a(std::move(x)), b(m, 0), tabl(build()) {} RMQpm1WithIndex(const std::vector<int>& x) : RMQpm1WithIndex(std::vector<int>(x)) {} std::pair<int, int> operator()(int l, int r) const { if (l >= r) return e(); static constexpr int MASK = SIZE - 1; auto f = [this](int l, int r) -> std::pair<int, int> { if (l >= r) return e(); int idx = cut(b[l >> LOG], l & MASK, ((r - 1) & MASK) + 1); return { a[l] + tabs.prd[idx], l + tabs.arg[idx] }; }; if (l >> LOG == (r - 1) >> LOG) return f(l, r); int spl = (l + SIZE - 1) >> LOG, spr = r >> LOG; return op(op(f(l, spl << LOG), f(spr << LOG, r)), tabl(spl, spr)); } private: int n, m; std::vector<int> a; std::vector<std::uint16_t> b; SparseTable<std::pair<int, int>, op, e> tabl; std::vector<std::pair<int, int>> build() { std::vector<std::pair<int, int>> c(m, e()); if (n == 0) return c; std::pair<int, int> p{ a[0] - 1, -1 }; for (int i = 0; i < n; p = { a[i], i }, ++i) { std::pair<int, int> q{ a[i], i }; int outer = i >> LOG; c[outer] = op(c[outer], q); b[outer] |= comp(q, p) << (i & (SIZE - 1)); } a.insert(a.begin(), a[0] - 1); assert(std::adjacent_find(a.begin(), a.end(), [](int x, int y) { return std::abs(x - y) != 1; }) == a.end()); return c; } static constexpr std::uint16_t cut(const std::uint16_t bits, const int l, const int r) { return std::uint16_t(bits << (SIZE - r)) >> (SIZE - r + l); } }; } // namespace suisen #line 5 "library/tree/lowest_common_ancestor.hpp" namespace suisen { class LowestCommonAncestor { public: LowestCommonAncestor(const std::vector<std::vector<int>> &g, int root = 0) : idx(g.size()), dep(2 * g.size() - 1), tour(2 * g.size() - 1), rmq(dfs(g, root)) {} int lca(int u, int v) const { return idx[u] <= idx[v] ? tour[rmq(idx[u], idx[v] + 1).second] : lca(v, u); } int dist(int u, int v) const { return dep[idx[u]] + dep[idx[v]] - 2 * dep[idx[operator()(u, v)]]; } int operator()(int u, int v) const { return lca(u, v); } private: std::vector<int> idx, dep, tour; RMQpm1WithIndex<true> rmq; std::vector<int>& dfs(const std::vector<std::vector<int>> &g, int root) { int k = 0; auto rec = [&](auto self, int u, int p, int d) -> void { dep[k] = d, tour[idx[u] = k++] = u; for (int v : g[u]) if (v != p) self(self, v, u, d + 1); if (p >= 0) { dep[k] = d - 1, tour[idx[p] = k++] = p; } }; rec(rec, root, -1, 0); return dep; } }; } // namespace suisen #line 1 "library/tree/level_ancestor.hpp" #line 6 "library/tree/level_ancestor.hpp" namespace suisen { struct LevelAncestor { LevelAncestor() = default; LevelAncestor(const std::vector<std::vector<int>>& g, const int root = 0) : _n(g.size()), _visit_time(_n), _visit_time_inv(_n), _depth(_n), _count(_n + 1), _bucket(_n) { build(g, root); } int query(const int u, const int k) const { if (k < 0 or k > _depth[u]) return -1; const int d = _depth[u] - k; const auto it_l = _bucket.begin() + _count[d], it_r = _bucket.begin() + _count[d + 1]; const int visit_time_u = _visit_time[u]; const int visit_time_v = *std::prev(std::upper_bound(it_l, it_r, visit_time_u)); return _visit_time_inv[visit_time_v]; } int operator()(const int u, const int k) const { return query(u, k); } int get_visit_time(const int u) const { return _visit_time[u]; } int get_visit_time_inv(const int t) const { return _visit_time_inv[t]; } int get_kth_visited(const int k) const { return _visit_time_inv[k]; } int get_depth(const int u) const { return _depth[u]; } private: int _n; std::vector<int> _visit_time; // vertex -> time std::vector<int> _visit_time_inv; // time -> vertex std::vector<int> _depth; // vertex -> depth std::vector<int> _count; // the leftmost index of the i'th block in the `_bucket` std::vector<int> _bucket; // [|dep=0|dep=1|dep=2|...|dep=n-1|]. Each block contains the visited times in ascending order. void build(const std::vector<std::vector<int>>& g, const int root) { int time = 0; auto dfs = [&](auto dfs, int u, int p) -> void { _visit_time[u] = time; _visit_time_inv[time] = u; ++time; ++_count[_depth[u] + 1]; for (int v : g[u]) if (v != p) { _depth[v] = _depth[u] + 1; dfs(dfs, v, u); } }; dfs(dfs, root, -1); for (int i = 0; i < _n; ++i) _count[i + 1] += _count[i]; auto index = _count; for (int v : _visit_time_inv) _bucket[index[_depth[v]]++] = _visit_time[v]; } }; } // namespace suisen #line 7 "test/src/tree/level_ancestor/jump_on_tree.test.cpp" int main() { std::ios::sync_with_stdio(false); std::cin.tie(nullptr); int n, q; std::cin >> n >> q; std::vector<std::vector<int>> g(n); for (int i = 0; i < n - 1; ++i) { int u, v; std::cin >> u >> v; g[u].push_back(v); g[v].push_back(u); } suisen::LowestCommonAncestor lca(g); suisen::LevelAncestor la(g); while (q --> 0) { int u, v, d; std::cin >> u >> v >> d; const int a = lca(u, v); const int du = la.get_depth(u) - la.get_depth(a); const int dv = la.get_depth(v) - la.get_depth(a); if (d <= du) { std::cout << la(u, d) << '\n'; } else if (d <= du + dv) { std::cout << la(v, du + dv - d) << '\n'; } else { std::cout << -1 << '\n'; } } return 0; }