test/yosupo/num_theory/stern_brocot_tree.test.cpp

• View this file on GitHub
• Last update: 2025-09-20 17:47:50+07:00
• Problem: https://judge.yosupo.jp/problem/stern_brocot_tree

Depends on

• Stern-Brocot Tree (math/stern_brocot_tree.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/stern_brocot_tree"

#include <bits/stdc++.h>
using namespace std;

#include "math/stern_brocot_tree.hpp"

using node = sbt_node<int>;

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int q;
    cin >> q;

    while (q--) {
        string op;
        cin >> op;
        if (op == "ENCODE_PATH") {
            int a, b;
            cin >> a >> b;
            auto path = node(a, b).path();
            cout << path.size();
            for (auto [d, n] : path)
                cout << ' ' << (!d ? 'L' : 'R') << ' ' << n;
            cout << '\n';
        } else if (op == "DECODE_PATH") {
            int k;
            cin >> k;
            node::sbt_path path(k);
            for (int i = 0; i < k; i++) {
                char d;
                int n;
                cin >> d >> n;
                path[i] = {d == 'R', n};
            }
            auto [a, b] = node(path).get();
            cout << a << ' ' << b << '\n';
        } else if (op == "LCA") {
            int a, b, c, d;
            cin >> a >> b >> c >> d;
            auto [f, g] = lca(node(a, b), node(c, d)).get();
            cout << f << ' ' << g << '\n';
        } else if (op == "ANCESTOR") {
            int k, a, b;
            cin >> k >> a >> b;
            auto x = node(a, b);
            if (x.go_up(x.depth() - k)) {
                tie(a, b) = x.get();
                cout << a << ' ' << b << '\n';
            } else {
                cout << "-1\n";
            }
        } else if (op == "RANGE") {
            int a, b;
            cin >> a >> b;
            auto [l, r] = node(a, b).range();
            auto [f, g] = l;
            auto [h, k] = r;
            cout << f << ' ' << g << ' ' << h << ' ' << k << '\n';
        }
    }
}

#line 1 "test/yosupo/num_theory/stern_brocot_tree.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/stern_brocot_tree"

#include <bits/stdc++.h>
using namespace std;

#line 2 "math/stern_brocot_tree.hpp"

// clang-format off
template<class Int> struct sbt_node {
    using node = sbt_node;
    using rational = pair<Int, Int>;
    using sbt_arc = bool;
    static constexpr sbt_arc Left = false, Right = true;
    using sbt_path = vector<pair<sbt_arc, Int>>;

    sbt_node() {}
    sbt_node(Int a, Int b) {
        assert(a > 0 && b > 0);
        sbt_arc dir = a < b ? Left : Right;
        if (dir == Left) swap(a, b);
        for (; b; dir = !dir) {
            const Int q = a / b, r = a % b;
            go_down(dir, q - (r == 0));
            a = exchange(b, r);
        }
    }
    sbt_node(const rational &x) : sbt_node(x.first, x.second) {}
    sbt_node(const sbt_path &path) {
        for (const auto &[dir, len] : path) go_down(dir, len);
    }

    operator rational() const { return {_l.first + _r.first, _l.second + _r.second}; }
    rational get() const { return *this; }
    pair<rational, rational> range() const { return {_l, _r}; }
    Int depth() const { return _dep; }
    const sbt_path &path() const { return _path; }

    friend bool operator==(const sbt_node &a, const sbt_node &b) { return a._l == b._l && a._r == b._r; }
    friend bool operator!=(const sbt_node &a, const sbt_node &b) { return !(a == b); }
    friend bool operator<(const sbt_node &a, const sbt_node &b) { return rational(a) < rational(b); }
    friend bool operator>(const sbt_node &a, const sbt_node &b) { return b > a; }
    friend bool operator<=(const sbt_node &a, const sbt_node &b) { return !(a > b); }
    friend bool operator>=(const sbt_node &a, const sbt_node &b) { return !(a < b); }

    friend sbt_node lca(const sbt_node &a, const sbt_node &b) {
        const sbt_path &pa = a.path(), &pb = b.path();
        const int k = (int)min(pa.size(), pb.size());
        sbt_node c;
        for (int i = 0; i < k; i++) {
            if (pa[i] == pb[i]) {
                c.go_down(pa[i].first, pa[i].second);
            } else {
                if (pa[i].first == pb[i].first)
                    c.go_down(pa[i].first, min(pa[i].second, pb[i].second));
                break;
            }
        }
        return c;
    }
    friend Int dist(const sbt_node &a, const sbt_node &b) {
        return a.depth() + b.depth() - 2 * lca(a, b).depth();
    }

    bool go_up(Int k) {
        if (k < 0 || depth() < k) return false;
        while (k) {
            auto &[dir, len] = _path.back();
            const Int u = min(k, len);
            k -= u;
            _dep -= u;
            if (dir == Left) {
                _r.first -= _l.first * u, _r.second -= _l.second * u;
            } else {
                _l.first -= _r.first * u, _l.second -= _r.second * u;
            }
            len -= u;
            if (len == 0) _path.pop_back();
        }
        return true;
    }
    void go_down(sbt_arc dir, Int k) {
        assert(k >= 0);
        if (k == 0) return;
        if (_path.size() && dir == _path.back().first) _path.back().second += k;
        else _path.emplace_back(dir, k);
        _dep += k;
        if (dir == Left) _r.first += _l.first * k, _r.second += _l.second * k;
        else _l.first += _r.first * k, _l.second += _r.second * k;
    }

  private:
    rational _l = {0, 1}, _r = {1, 0};
    Int _dep = 0;
    sbt_path _path;

    static bool less(const rational &a, const rational &b) {
        using LInt = conditional_t<numeric_limits<Int>::digits <= 32, uint64_t,
                                   __uint128_t>;
        return LInt(a.first) * b.second < LInt(b.first) * a.second;
    }
};
// clang-format on
#line 7 "test/yosupo/num_theory/stern_brocot_tree.test.cpp"

using node = sbt_node<int>;

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int q;
    cin >> q;

    while (q--) {
        string op;
        cin >> op;
        if (op == "ENCODE_PATH") {
            int a, b;
            cin >> a >> b;
            auto path = node(a, b).path();
            cout << path.size();
            for (auto [d, n] : path)
                cout << ' ' << (!d ? 'L' : 'R') << ' ' << n;
            cout << '\n';
        } else if (op == "DECODE_PATH") {
            int k;
            cin >> k;
            node::sbt_path path(k);
            for (int i = 0; i < k; i++) {
                char d;
                int n;
                cin >> d >> n;
                path[i] = {d == 'R', n};
            }
            auto [a, b] = node(path).get();
            cout << a << ' ' << b << '\n';
        } else if (op == "LCA") {
            int a, b, c, d;
            cin >> a >> b >> c >> d;
            auto [f, g] = lca(node(a, b), node(c, d)).get();
            cout << f << ' ' << g << '\n';
        } else if (op == "ANCESTOR") {
            int k, a, b;
            cin >> k >> a >> b;
            auto x = node(a, b);
            if (x.go_up(x.depth() - k)) {
                tie(a, b) = x.get();
                cout << a << ' ' << b << '\n';
            } else {
                cout << "-1\n";
            }
        } else if (op == "RANGE") {
            int a, b;
            cin >> a >> b;
            auto [l, r] = node(a, b).range();
            auto [f, g] = l;
            auto [h, k] = r;
            cout << f << ' ' << g << ' ' << h << ' ' << k << '\n';
        }
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	edge_decode_00	AC	42 ms	3 MB
g++	edge_encode_00	AC	46 ms	3 MB
g++	edge_lca_00	AC	30 ms	3 MB
g++	edge_range_00	AC	34 ms	4 MB
g++	example_00	AC	5 ms	4 MB
g++	hand_00	AC	4 ms	4 MB
g++	hand_01	AC	4 ms	4 MB
g++	random_ancestor_no_00	AC	42 ms	4 MB
g++	random_ancestor_yes_00	AC	55 ms	4 MB
g++	random_decode_00	AC	125 ms	4 MB
g++	random_encode_00	AC	137 ms	4 MB
g++	random_lca_00	AC	78 ms	4 MB
g++	random_range_00	AC	58 ms	4 MB
g++	small_ancestor_00	AC	28 ms	4 MB
g++	small_decode_00	AC	55 ms	4 MB
g++	small_encode_00	AC	61 ms	4 MB
g++	small_lca_00	AC	37 ms	3 MB
g++	small_range_00	AC	40 ms	4 MB
clang++	edge_decode_00	AC	41 ms	4 MB
clang++	edge_encode_00	AC	47 ms	4 MB
clang++	edge_lca_00	AC	28 ms	4 MB
clang++	edge_range_00	AC	28 ms	4 MB
clang++	example_00	AC	4 ms	4 MB
clang++	hand_00	AC	4 ms	4 MB
clang++	hand_01	AC	4 ms	4 MB
clang++	random_ancestor_no_00	AC	43 ms	4 MB
clang++	random_ancestor_yes_00	AC	55 ms	4 MB
clang++	random_decode_00	AC	121 ms	4 MB
clang++	random_encode_00	AC	156 ms	4 MB
clang++	random_lca_00	AC	75 ms	4 MB
clang++	random_range_00	AC	55 ms	4 MB
clang++	small_ancestor_00	AC	28 ms	4 MB
clang++	small_decode_00	AC	57 ms	4 MB
clang++	small_encode_00	AC	59 ms	4 MB
clang++	small_lca_00	AC	35 ms	4 MB
clang++	small_range_00	AC	40 ms	4 MB

Back to top page