Test/LC/stern_brocot_tree.test.cpp

View this file on GitHub
Last update: 2025-08-23 22:47:55+09:00
Problem: https://judge.yosupo.jp/problem/stern_brocot_tree

Depends on

Code

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

#include "../../Src/Number/SternBrocotTree.hpp"

#include <cassert>
#include <iostream>
#include <vector>
#include <string>
using namespace std;
using namespace zawa;
using Node = SternBrocotTree::Node<int>;
using Path = SternBrocotTree::Path<int>;
char from(SternBrocotTree::Direction d) {
    return d == SternBrocotTree::Left ? 'L' : 'R';
}
SternBrocotTree::Direction to(char c) {
    assert(c == 'L' or c == 'R');
    return c == 'L' ? SternBrocotTree::Left : SternBrocotTree::Right;
}
int main() {
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    int Q;
    cin >> Q;
    while (Q--) {
        string T;
        cin >> T;
        if (T == "ENCODE_PATH") {
            int a, b;
            cin >> a >> b;
            Path ans = SternBrocotTree::Encode(Node{a, b});
            cout << ssize(ans);
            for (auto [d, s] : ans)
                cout << ' ' << from(d) << ' ' << s;
            cout << '\n';
        }
        else if (T == "DECODE_PATH") {
            int k;
            cin >> k;
            Path path(k);
            for (int i = 0 ; i < k ; i++) {
                char c;
                int s;
                cin >> c >> s;
                path[i] = {to(c), s};
            }
            auto [a, b] = SternBrocotTree::Decode(path);
            cout << a << ' ' << b << '\n';
        }
        else if (T == "LCA") {
            int a, b, c, d;
            std::cin >> a >> b >> c >> d;
            auto [p, q] = SternBrocotTree::LCA(Node{a, b}, Node{c, d});
            cout << p << ' ' << q << '\n';
        }
        else if (T == "ANCESTOR") {
            int k, a, b;
            cin >> k >> a >> b;
            auto ans = SternBrocotTree::Ancestor(k, Node{a, b});
            if (ans.has_value())
                cout << ans->first << ' ' << ans->second << '\n';
            else
                cout << -1 << '\n';
        }
        else if (T == "RANGE") {
            int a, b;
            cin >> a >> b;
            auto [l, r] = SternBrocotTree::Range(Node{a, b});
            cout << l.first << ' ' << l.second << ' ' << r.first << ' ' << r.second << '\n';
        }
        else
            assert(0);
    }
}

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

#line 2 "Src/Number/SternBrocotTree.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Number/SternBrocotTree.hpp"

#include <algorithm>
#include <concepts>
#include <optional>
#include <utility>
#include <vector>
#include <tuple>

namespace zawa {

namespace SternBrocotTree {

namespace sbt_internal {

template <class T>
concept Integer = std::integral<T> or std::same_as<T, __int128_t> or std::same_as<T, __uint128_t>;

} // namespace sbt_internal

enum Direction {
    Left,
    Right
};

template <sbt_internal::Integer T>
using Path = std::vector<std::pair<Direction, T>>;

template <sbt_internal::Integer T>
using Node = std::pair<T, T>;

template <sbt_internal::Integer T>
Path<T> Encode(Node<T> node) {
    auto [a, b] = node;
    u32 d = [&]() {
        if (a < b) {
            std::swap(a, b);
            return 0;
        }
        else {
            return 1;
        }
    }();
    std::vector<T> quos;
    while (b) {
        quos.push_back(a / b);
        T r = a % b;
        a = b;
        b = r;
    }
    if (quos.size()) {
        quos.back()--;
        if (quos.back() == 0)
            quos.pop_back();
    }
    Path<T> res;
    for (T q : quos) {
        res.push_back({static_cast<Direction>(d), q});
        d ^= 1;
    }
    return res;
}

namespace sbt_internal {

template <Integer T>
std::tuple<Node<T>, Node<T>, Node<T>> Decode(const Path<T>& path) {
    T a = 0, b = 1, c = 1, d = 1, e = 1, f = 0;
    for (auto [direction, step] : path) {
        if (direction == Left) {
            e = c + (step - 1) * a;
            f = d + (step - 1) * b;
            c += step * a;
            d += step * b;
        }
        else if (direction == Right) {
            a = c + (step - 1) * e;
            b = d + (step - 1) * f;
            c += step * e;
            d += step * f;
        } 
    }
    return {Node<T>{a, b}, Node<T>{c, d}, Node<T>{e, f}};
}

} // namespace sbt_internal

template <sbt_internal::Integer T>
Node<T> Decode(const Path<T>& path) {
    return std::get<1>(sbt_internal::Decode(path));
}

template <sbt_internal::Integer T>
Node<T> LCA(Node<T> p, Node<T> q) {
    auto a = Encode(p), b = Encode(q);
    Path<T> path;
    for (usize i = 0 ; i < std::min(a.size(), b.size()) ; i++) {
        if (a[i].first != b[i].first)
            break;
        path.push_back({a[i].first, std::min(a[i].second, b[i].second)});
        if (a[i].second != b[i].second)
            break;
    }
    return Decode(path);
}

// depth is 0-indexed
template <sbt_internal::Integer T>
std::optional<Node<T>> Ancestor(u64 dep, Node<T> n) {
    Path<T> path;
    for (auto [direction, step] : Encode(n)) {
        if (dep == 0)
           break;
        u64 s = std::min<u64>(dep, step);
        path.push_back({direction, s});
        dep -= s;
    }
    if (dep)
        return std::nullopt;
    else
        return Decode(path);
}

template <sbt_internal::Integer T>
std::pair<Node<T>, Node<T>> Range(Node<T> n) {
    Node<T> l, r;
    std::tie(l, std::ignore, r) = sbt_internal::Decode(Encode(n));
    return {l, r};
}

} // namespace SternBrocotTree


} // namespace zawa
#line 4 "Test/LC/stern_brocot_tree.test.cpp"

#include <cassert>
#include <iostream>
#line 8 "Test/LC/stern_brocot_tree.test.cpp"
#include <string>
using namespace std;
using namespace zawa;
using Node = SternBrocotTree::Node<int>;
using Path = SternBrocotTree::Path<int>;
char from(SternBrocotTree::Direction d) {
    return d == SternBrocotTree::Left ? 'L' : 'R';
}
SternBrocotTree::Direction to(char c) {
    assert(c == 'L' or c == 'R');
    return c == 'L' ? SternBrocotTree::Left : SternBrocotTree::Right;
}
int main() {
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    int Q;
    cin >> Q;
    while (Q--) {
        string T;
        cin >> T;
        if (T == "ENCODE_PATH") {
            int a, b;
            cin >> a >> b;
            Path ans = SternBrocotTree::Encode(Node{a, b});
            cout << ssize(ans);
            for (auto [d, s] : ans)
                cout << ' ' << from(d) << ' ' << s;
            cout << '\n';
        }
        else if (T == "DECODE_PATH") {
            int k;
            cin >> k;
            Path path(k);
            for (int i = 0 ; i < k ; i++) {
                char c;
                int s;
                cin >> c >> s;
                path[i] = {to(c), s};
            }
            auto [a, b] = SternBrocotTree::Decode(path);
            cout << a << ' ' << b << '\n';
        }
        else if (T == "LCA") {
            int a, b, c, d;
            std::cin >> a >> b >> c >> d;
            auto [p, q] = SternBrocotTree::LCA(Node{a, b}, Node{c, d});
            cout << p << ' ' << q << '\n';
        }
        else if (T == "ANCESTOR") {
            int k, a, b;
            cin >> k >> a >> b;
            auto ans = SternBrocotTree::Ancestor(k, Node{a, b});
            if (ans.has_value())
                cout << ans->first << ' ' << ans->second << '\n';
            else
                cout << -1 << '\n';
        }
        else if (T == "RANGE") {
            int a, b;
            cin >> a >> b;
            auto [l, r] = SternBrocotTree::Range(Node{a, b});
            cout << l.first << ' ' << l.second << ' ' << r.first << ' ' << r.second << '\n';
        }
        else
            assert(0);
    }
}