cp-library-cpp

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub suisen-cp/cp-library-cpp

:heavy_check_mark: test/src/tree/level_ancestor/jump_on_tree.test.cpp

Depends on

Code

#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;
}
Back to top page