test/tree/jump_on_tree.0.test.cpp

• View this file on GitHub
• Last update: 2026-02-11 21:03:28+08:00
• Problem: https://judge.yosupo.jp/problem/jump_on_tree

Depends on

• node_pool.hpp
• st_tree_node_base.hpp

Code

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

#include "node_pool.hpp"
#include "st_tree_node_base.hpp"
#include <iostream>

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int n, q;
    std::cin >> n >> q;
    struct STTreeNode : STTreeNodeBase<STTreeNode> {};
    FixedSizeNodePool<STTreeNode> pool(n);
    auto [node, id] = pool.get_func();
    for (int i = 0; i < n - 1; ++i) {
        int a, b;
        std::cin >> a >> b;
        node(a)->link(node(b));
    }
    while (q--) {
        int s, t, i;
        std::cin >> s >> t >> i;
        node(t)->evert();
        node(s)->expose();
        if (node(s)->size() > i) {
            std::cout << id(node(s)->jump(i)) << '\n';
        } else {
            std::cout << "-1\n";
        }
    }
    return 0;
}

#line 1 "test/tree/jump_on_tree.0.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/jump_on_tree

#line 2 "node_pool.hpp"

#include <list>
#include <memory>
#include <utility>
#include <vector>

template<typename NodeT> class FixedSizeNodePool {
    std::vector<NodeT> pool;

public:
    explicit FixedSizeNodePool(int n) : pool(n) {}
    NodeT *at(int ind) { return pool.data() + ind; }
    int id(NodeT *a) const { return a - pool.data(); }
    auto get_func() {
        return std::make_pair([this](int ind) { return at(ind); },
                              [this](NodeT *a) { return id(a); });
    }
};

template<typename NodeT> class DynamicSizeNodePool {
    struct Wrapped : public NodeT {
        using NodeT::NodeT;
        typename std::list<Wrapped>::iterator i;
    };
    std::list<Wrapped> used_list;
    std::list<Wrapped> free_list;

public:
    template<typename... Args> NodeT *make(Args &&...arg) {
        if (free_list.empty()) {
            auto &&node = used_list.emplace_back(std::forward<Args>(arg)...);
            node.i      = std::prev(used_list.end());
            return std::addressof(node);
        }
        used_list.splice(used_list.end(), free_list, std::prev(free_list.end()));
        auto &&node = used_list.back();
        node.~NodeT(); // i remains unchanged
        new (std::addressof(node)) NodeT(std::forward<Args>(arg)...);
        return std::addressof(node);
    }
    // this is lazy, if sth. relies on the order of dtor, do NOT use
    void retrieve(NodeT *node) {
        free_list.splice(free_list.end(), used_list, ((Wrapped *)node)->i);
    }
};
#line 2 "st_tree_node_base.hpp"

#line 4 "st_tree_node_base.hpp"

template<typename STTreeNodeT> class STTreeNodeBase {
    STTreeNodeBase *L;
    STTreeNodeBase *R;
    STTreeNodeBase *P;
    int Size;
    bool NeedFlip;

    enum class Child { LEFT, RIGHT };
    Child which() const { return P->L == this ? Child::LEFT : Child::RIGHT; }
    // has NO parent OR NOT a prefered child
    bool is_root() const { return P == nullptr || (P->L != this && P->R != this); }
    bool is_left_child() const { return which() == Child::LEFT; }
    bool is_right_child() const { return which() == Child::RIGHT; }

    // CRTP reimplement
    void do_flip() {}
    void do_propagate() {}
    void do_update() {}

protected:
    void base_flip() {
        NeedFlip = !NeedFlip;
        std::swap(L, R);
        ((STTreeNodeT &)*this).do_flip();
    }
    // base_propagate() is called to propagate the update information to child(ren).
    // There is no need to update the information combined from child(ren)
    // which should be done in base_update().
    void base_propagate() {
        ((STTreeNodeT &)*this).do_propagate();
        if (NeedFlip) {
            NeedFlip = false;
            if (L) L->base_flip();
            if (R) R->base_flip();
        }
    }
    // base_update() is called to update the information combined from child(ren).
    void base_update() {
        Size = 1;
        if (L) Size += L->Size;
        if (R) Size += R->Size;
        ((STTreeNodeT &)*this).do_update();
    }
    void base_rotate() {
        P->base_propagate();
        base_propagate();
        if (is_left_child()) {
            if ((P->L = R)) R->P = P;
            if (!P->is_root()) {
                if (P->is_left_child()) P->P->L = this;
                else { P->P->R = this; }
            }
            R = P, P = P->P, R->P = this;
            R->base_update();
        } else {
            if ((P->R = L)) L->P = P;
            if (!P->is_root()) {
                if (P->is_left_child()) P->P->L = this;
                else { P->P->R = this; }
            }
            L = P, P = P->P, L->P = this;
            L->base_update();
        }
    }
    void base_splay() {
        for (base_propagate(); !is_root(); base_rotate()) {
            if (!P->is_root()) {
                P->P->base_propagate();
                P->which() == which() ? P->base_rotate() : base_rotate();
            }
        }
        base_update();
    }

    STTreeNodeBase() : L(), R(), P(), Size(1), NeedFlip() {}

public:
    int size() const { return Size; }

    STTreeNodeT *left() const { return (STTreeNodeT *)L; }
    STTreeNodeT *right() const { return (STTreeNodeT *)R; }

    void update() { base_update(); }

    STTreeNodeT *expose() {
        STTreeNodeBase *a = this, *lca = a;
        base_splay();
        a->R = nullptr;
        while (a->P) {
            lca = a->P;
            lca->base_splay();
            a->P->R = a;
            a->base_rotate();
        }
        a->base_update();
        // now a is the root of the aux. tree
        return (STTreeNodeT *)lca;
    }
    void evert() { expose(), base_flip(); }
    STTreeNodeT *root() {
        expose();
        STTreeNodeBase *a = this;
        while (a->L) a = a->L;
        a->base_splay();
        return (STTreeNodeT *)a;
    }
    STTreeNodeT *parent() {
        expose();
        if (!L) return nullptr;
        STTreeNodeBase *a = L;
        a->base_propagate();
        while (a->R) {
            a = a->R;
            a->base_propagate();
        }
        a->base_splay();
        return (STTreeNodeT *)a;
    }
    // this op. WILL change the root
    void link(STTreeNodeT *a) {
        evert();
        if (a->root() != (STTreeNodeT *)this) P = a;
    }
    // this op. will NOT change the root
    void cut() {
        expose();
        STTreeNodeBase *b = L;
        L                 = nullptr;
        if (b) b->P = nullptr;
        base_update();
    }
    // this op. will NOT change the root
    void cut(STTreeNodeT *b) {
        if (parent() == b) {
            cut();
        } else if (b->parent() == (STTreeNodeT *)this) {
            b->cut();
        }
    }
    STTreeNodeT *select(int k) {
        // example: A --> B --> C --> D, where D is the root of the path
        // We use expose(A) than:
        // * size(A) = 4
        // * select(A, 0) returns D
        // * select(A, 1) returns C
        // * ...
        STTreeNodeBase *a = this;
        a->base_propagate();
        while ((a->L ? a->L->Size : 0) != k) {
            if ((a->L ? a->L->Size : 0) < k) {
                k -= (a->L ? a->L->Size : 0) + 1;
                a = a->R;
            } else {
                a = a->L;
            }
            a->base_propagate();
        }
        a->base_splay();
        return (STTreeNodeT *)a;
    }
    STTreeNodeT *jump(int k) { return select(Size - k - 1); }
};
#line 5 "test/tree/jump_on_tree.0.test.cpp"
#include <iostream>

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int n, q;
    std::cin >> n >> q;
    struct STTreeNode : STTreeNodeBase<STTreeNode> {};
    FixedSizeNodePool<STTreeNode> pool(n);
    auto [node, id] = pool.get_func();
    for (int i = 0; i < n - 1; ++i) {
        int a, b;
        std::cin >> a >> b;
        node(a)->link(node(b));
    }
    while (q--) {
        int s, t, i;
        std::cin >> s >> t >> i;
        node(t)->evert();
        node(s)->expose();
        if (node(s)->size() > i) {
            std::cout << id(node(s)->jump(i)) << '\n';
        } else {
            std::cout << "-1\n";
        }
    }
    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
clang++	almost_line_00	AC	2626 ms	40 MB
clang++	almost_line_01	AC	2388 ms	40 MB
clang++	almost_line_02	AC	2400 ms	36 MB
clang++	almost_line_03	AC	2382 ms	38 MB
clang++	almost_uni_00	AC	385 ms	35 MB
clang++	almost_uni_01	AC	383 ms	37 MB
clang++	example_00	AC	9 ms	20 MB
clang++	line_00	AC	2241 ms	37 MB
clang++	max_random_00	AC	1286 ms	40 MB
clang++	max_random_01	AC	1283 ms	37 MB
clang++	random_00	AC	1213 ms	34 MB
clang++	random_01	AC	1243 ms	34 MB
clang++	random_02	AC	859 ms	25 MB
clang++	random_03	AC	1294 ms	37 MB
clang++	random_04	AC	1113 ms	28 MB
clang++	small_random_00	AC	283 ms	22 MB
clang++	small_random_01	AC	287 ms	20 MB
clang++	small_random_02	AC	334 ms	16 MB
clang++	small_random_03	AC	315 ms	24 MB
clang++	small_random_04	AC	257 ms	23 MB
clang++	uni_00	AC	347 ms	37 MB

Back to top page