Test/DataStructure/SternBrocotTree.test.cpp

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

Depends on

• Library/DataStructure/SternBrocotTree.hpp
• Library/Range/util.hpp
• Library/Utility/Tuple.hpp

Code

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

#include <iostream>

#include <ranges>


// begin:tag includes

#include "./../../Library/DataStructure/SternBrocotTree.hpp"

#include "./../../Library/Range/util.hpp"

// end:tag includes


using ll = long long;

signed main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(0);

  int t;
  std::cin >> t;
  mtd::SternBrocotTree<ll> sbt;
  for (auto _ : std::views::iota(0, t)) {
    std::string s;
    std::cin >> s;

    if (s == "ENCODE_PATH") {
      ll a, b;
      std::cin >> a >> b;
      auto path_rle = sbt.encode(a, b);
      std::cout << path_rle.size() << (path_rle.empty() ? "" : " ");
      for (const auto& [i, right, k] :
           path_rle | std::views::enumerate | mtd::views::flatten) {
        std::cout << (right ? 'R' : 'L') << " " << k
                  << (i == path_rle.size() - 1 ? "" : " ");
      }
      std::cout << std::endl;
    }

    if (s == "DECODE_PATH") {
      ll k;
      std::cin >> k;
      std::vector<std::tuple<bool, ll>> path_rle;
      for (auto __ : std::views::iota(0, k)) {
        char c;
        ll n;
        std::cin >> c >> n;
        path_rle.emplace_back(c == 'R', n);
      }
      auto [a, b] = sbt.decode(path_rle).get();
      std::cout << a << " " << b << std::endl;
    }

    if (s == "LCA") {
      ll a, b, c, d;
      std::cin >> a >> b >> c >> d;
      auto [f, g] = sbt.lca(a, b, c, d).get();
      std::cout << f << " " << g << std::endl;
    }

    if (s == "ANCESTOR") {
      ll k, a, b;
      std::cin >> k >> a >> b;
      try {
        auto [f, g] = sbt.ancestor(a, b, k).get();
        std::cout << f << " " << g << std::endl;
      } catch (const std::runtime_error& e) { std::cout << -1 << std::endl; }
    }

    if (s == "RANGE") {
      ll a, b;
      std::cin >> a >> b;
      try {
        auto [node_l, node_r] = sbt.range(a, b);
        auto [f, g] = node_l.get();
        auto [h, k] = node_r.get();
        std::cout << f << " " << g << " " << h << " " << k << std::endl;
      } catch (const std::runtime_error& e) {
        std::cout << e.what() << std::endl;
      }
    }
  }
}

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

#include <iostream>

#include <ranges>


// begin:tag includes

#line 2 "Library/DataStructure/SternBrocotTree.hpp"

#line 4 "Library/DataStructure/SternBrocotTree.hpp"
#include <numeric>
#line 6 "Library/DataStructure/SternBrocotTree.hpp"
#include <stdexcept>
#include <tuple>
#include <vector>

namespace mtd {

  template <class T, class CompT = long long>
  class SternBrocotTree {
    using Path = std::vector<std::tuple<bool, T>>;

    static constexpr T MAX_NUM = static_cast<T>(2e18);
    static constexpr T MAX_DEN = static_cast<T>(2e18);

    class Node {
      // 定数倍高速化のため破壊的変更や怪しい仕様あり
      T num_l, den_l, num_r, den_r;
      Path path_rle;
      const T max_num;
      const T max_den;

      friend std::ostream& operator<<(std::ostream& os, const Node& node) {
        return os << node.num_l + node.num_r << "/" << node.den_l + node.den_r
                  << ": " << node.num_l << "/" << node.den_l << " "
                  << node.num_r << "/" << node.den_r;
      }

    public:
      static constexpr auto get_root(T max_num, T max_den) {
        return Node(0, 1, 1, 0, Path(), max_num, max_den);
      }

      constexpr auto get() const {
        return std::make_tuple(num_l + num_r, den_l + den_r);
      }
      constexpr auto get_l() const { return Node(num_l, den_l); }
      constexpr auto get_r() const { return Node(num_r, den_r); }
      constexpr auto get_path_rle() const { return path_rle; }

      constexpr auto move_left(T d = 1) {
        if (num_l > 0) { d = std::min(d, (max_num - num_r - num_l) / num_l); }
        if (den_l > 0) { d = std::min(d, (max_den - den_r - den_l) / den_l); }
        if (d <= 0) { return false; }
        path_rle.emplace_back(false, d);
        num_r += d * num_l;
        den_r += d * den_l;
        return true;
      }
      constexpr auto move_left_to(T num, T den) {
        auto den_d = static_cast<CompT>(den);
        auto num_d = static_cast<CompT>(num);
        auto tmp = den_l * num_d - den_d * num_l;
        T d =
            (den_d * (num_l + num_r) - (den_l + den_r) * num_d + tmp - 1) / tmp;
        return move_left(d);
      }
      constexpr auto move_right(T d = 1) {
        if (num_r > 0) { d = std::min(d, (max_num - num_l - num_r) / num_r); }
        if (den_r > 0) { d = std::min(d, (max_den - den_l - den_r) / den_r); }
        if (d <= 0) { return false; }
        path_rle.emplace_back(true, d);
        num_l += d * num_r;
        den_l += d * den_r;
        return true;
      }
      constexpr auto move_right_to(T num, T den) {
        auto den_d = static_cast<CompT>(den);
        auto num_d = static_cast<CompT>(num);
        auto tmp = den_d * num_r - den_r * num_d;
        T d =
            ((den_l + den_r) * num_d - den_d * (num_l + num_r) + tmp - 1) / tmp;
        return move_right(d);
      }

      constexpr static auto generate_node(T num, T den, T max_num, T max_den) {
        if (den <= 0) {
          throw std::runtime_error("denominator must be positive");
        }
        if (num < 0) {
          throw std::runtime_error("numerator must be non-negative");
        }
        if (std::gcd(num, den) > 1) {
          throw std::runtime_error("numerator and denominator must be coprime");
        }

        auto node = get_root(max_num, max_den);
        while (node.move_left_to(num, den) || node.move_right_to(num, den)) {}
        return node;
      }

      constexpr static auto decode(const Path& path_rle, T max_num = MAX_NUM,
                                   T max_den = MAX_DEN) {
        auto node = get_root(max_num, max_den);
        for (const auto& [right, k] : path_rle) {
          right ? node.move_right(k) : node.move_left(k);
        }
        return node;
      }

      constexpr Node(T num_l, T den_l, T num_r, T den_r, Path&& path_rle,
                     T max_num, T max_den)
          : num_l(num_l),
            den_l(den_l),
            num_r(num_r),
            den_r(den_r),
            path_rle(std::move(path_rle)),
            max_num(max_num),
            max_den(max_den) {}
      constexpr Node(T num_l, T den_l, T num_r, T den_r)
          : Node(num_l, den_l, num_r, den_r, Path(), MAX_NUM, MAX_DEN) {}
      constexpr Node(T num, T den)
          : Node(generate_node(num, den, MAX_NUM, MAX_DEN)) {}

      constexpr auto operator!=(const Node& other) const {
        return std::tie(num_l, den_l, num_r, den_r) !=
               std::tie(other.num_l, other.den_l, other.num_r, other.den_r);
      }
      constexpr auto operator==(const Node& other) const {
        return !(*this != other);
      }
    };

  public:
    /*
     * Encode the path from the root to the fraction num/den
     **/
    constexpr auto encode(const Node& node) const {
      return node.get_path_rle();
    }
    constexpr auto encode(T num, T den) const { return encode(Node(num, den)); }

    /*
     * Decode the path from the root to the fraction represented by
     **/
    constexpr auto decode(const Path& path_rle) const {
      return Node::decode(path_rle);
    }

    /*
     * Find the lowest common ancestor of two fractions num1/den1 and num2/den2
     **/
    constexpr auto lca(const Node& node1, const Node& node2) const {
      auto path_rle1 = encode(node1);
      auto path_rle2 = encode(node2);
      Path lca_path;
      for (const auto [p1, p2] : std::views::zip(path_rle1, path_rle2)) {
        auto [right1, k1] = p1;
        auto [right2, k2] = p2;
        if (right1 != right2) { return Node::get_root(MAX_NUM, MAX_DEN); }
        lca_path.emplace_back(right1, std::min(k1, k2));
        if (p1 != p2) { break; }
      }
      return decode(lca_path);
    }
    constexpr auto lca(T num1, T den1, T num2, T den2) const {
      return lca(Node(num1, den1), Node(num2, den2));
    }

    /*
     * Find the k-th ancestor of the fraction num/den
     **/
    constexpr auto ancestor(const Node& node, T k) const {
      Path k_path_rle;
      for (const auto& [right, count] : encode(node)) {
        if (count > k) {
          k_path_rle.emplace_back(right, k);
          k = 0;
          break;
        } else {
          k_path_rle.emplace_back(right, count);
          k -= count;
        }
      }
      if (k > 0) { throw std::runtime_error("k is too large for the path"); }
      return decode(k_path_rle);
    }
    constexpr auto ancestor(T num, T den, T k) const {
      return ancestor(Node(num, den), k);
    }

    /*
     * Find the lower and upper bounds of the descendants of num/den
     **/
    constexpr auto range(const Node& node) const {
      auto [num, den] = node.get();
      if (num == 1 && den == 1) {
        return std::make_tuple(Node(0, 0, 0, 1), Node(0, 0, 1, 0));
      }
      if (den == 1) { return std::make_tuple(node.get_l(), Node(0, 0, 1, 0)); }
      if (num == 1) { return std::make_tuple(Node(0, 0, 0, 1), node.get_r()); }
      return std::make_tuple(node.get_l(), node.get_r());
    }
    constexpr auto range(T num, T den) const { return range(Node(num, den)); }

    /*
     * Create a node representing the fraction num/den
     **/
    constexpr auto create_node(T num, T den, T max_num = MAX_NUM,
                               T max_den = MAX_DEN) const {
      return Node::generate_node(num, den, max_num, max_den);
    }

    /*
     * Get the root node of the tree
     **/
    constexpr auto get_root(T max_num = MAX_NUM, T max_den = MAX_DEN) const {
      return Node::get_root(max_num, max_den);
    }
  };
}  // namespace mtd
#line 2 "Library/Range/util.hpp"

#include <algorithm>
#line 6 "Library/Range/util.hpp"

#line 2 "Library/Utility/Tuple.hpp"

#include <functional>

namespace mtd {
  namespace util {
    template <class F, class T>
    constexpr auto tuple_transform(F&& f, T&& t) {
      return std::apply(
          [&]<class... Ts>(Ts&&... elems) {
            return std::tuple<std::invoke_result_t<F&, Ts>...>(
                std::invoke(f, std::forward<Ts>(elems))...);
          },
          std::forward<T>(t));
    }
    template <class F, class T>
    constexpr auto tuple_for_each(F&& f, T&& t) {
      std::apply(
          [&]<class... Ts>(Ts&&... elems) {
            (std::invoke(f, std::forward<Ts>(elems)), ...);
          },
          std::forward<T>(t));
    }
  }  // namespace util
}  // namespace mtd
#line 8 "Library/Range/util.hpp"

namespace mtd {
  namespace ranges {

    namespace __detail {
      template <typename... T>
      concept __all_random_access = (std::ranges::random_access_range<T> &&
                                     ...);
      template <typename... T>
      concept __all_bidirectional = (std::ranges::bidirectional_range<T> &&
                                     ...);
      template <typename... T>
      concept __all_forward = (std::ranges::forward_range<T> && ...);

      template <class... T>
      constexpr auto _S_iter_concept() {
        if constexpr (__all_random_access<T...>) {
          return std::random_access_iterator_tag{};
        } else if constexpr (__all_bidirectional<T...>) {
          return std::bidirectional_iterator_tag{};
        } else if constexpr (__all_forward<T...>) {
          return std::forward_iterator_tag{};
        } else {
          return std::input_iterator_tag{};
        }
      }

      template <typename T>
      auto _flatten(const T& t) {
        return std::make_tuple(t);
      }
      template <typename... T>
      auto _flatten(const std::tuple<T...>& t);

      template <typename Head, typename... Tail>
      auto _flatten_impl(const Head& head, const Tail&... tail) {
        return std::tuple_cat(_flatten(head), _flatten(tail)...);
      }
      template <typename... T>
      auto _flatten(const std::tuple<T...>& t) {
        return std::apply(
            [](const auto&... args) { return _flatten_impl(args...); }, t);
      }
    }  // namespace __detail

    template <std::ranges::range _Range>
    struct flatten_view
        : public std::ranges::view_interface<flatten_view<_Range>> {
      class iterator {
      public:
        std::ranges::iterator_t<_Range> _M_current;

        using difference_type = std::ranges::range_difference_t<_Range>;
        using value_type = decltype(__detail::_flatten(
            std::declval<
                std::iter_reference_t<std::ranges::iterator_t<_Range>>>()));
        using iterator_concept = decltype(__detail::_S_iter_concept<_Range>());

        constexpr iterator() = default;
        constexpr explicit iterator(decltype(_M_current) __current)
            : _M_current(__current) {}
        constexpr auto operator*() const {
          return __detail::_flatten(*_M_current);
        }
        constexpr auto& operator++() {
          ++_M_current;
          return *this;
        }
        constexpr auto operator++(int) { return ++*this; }
        constexpr auto operator==(const iterator& other) const {
          return _M_current == other._M_current;
        }
        constexpr auto& operator--() requires
            __detail::__all_bidirectional<_Range> {
          --_M_current;
          return *this;
        }
        constexpr auto operator--(
            int) requires __detail::__all_bidirectional<_Range> {
          return --*this;
        }
        constexpr auto operator<=>(const iterator&)
            const requires __detail::__all_random_access<_Range>
        = default;
        constexpr auto operator-(const iterator& itr)
            const requires __detail::__all_random_access<_Range> {
          return _M_current - itr._M_current;
        }
        constexpr auto& operator+=(const difference_type n) requires
            __detail::__all_random_access<_Range> {
          _M_current += n;
          return *this;
        }
        constexpr auto operator+(const difference_type n)
            const requires __detail::__all_random_access<_Range> {
          auto __tmp = *this;
          __tmp += n;
          return __tmp;
        }
        constexpr friend auto operator+(const difference_type n,
                                        const iterator& itr) requires
            __detail::__all_random_access<_Range> {
          return itr + n;
        }
        constexpr auto& operator-=(const difference_type n) requires
            __detail::__all_random_access<_Range> {
          _M_current -= n;
          return *this;
        }
        constexpr auto operator-(const difference_type n)
            const requires __detail::__all_random_access<_Range> {
          auto __tmp = *this;
          __tmp -= n;
          return __tmp;
        }
        constexpr auto operator[](const difference_type n)
            const requires __detail::__all_random_access<_Range> {
          return __detail::_flatten(_M_current[n]);
        }
      };

      class sentinel {
        std::ranges::sentinel_t<_Range> _M_end;

      public:
        constexpr sentinel() = default;
        constexpr explicit sentinel(const decltype(_M_end)& __end)
            : _M_end(__end) {}

        friend constexpr bool operator==(const iterator& __x,
                                         const sentinel& __y) {
          return __x._M_current == __y._M_end;
        }
      };

      _Range _M_views;
      constexpr explicit flatten_view(const _Range& __views)
          : _M_views(__views) {}
      constexpr auto begin() { return iterator(std::ranges::begin(_M_views)); }
      constexpr auto end() { return sentinel(std::ranges::end(_M_views)); }
    };

  }  // namespace ranges

  namespace views {
    namespace __detail {
      template <typename... _Args>
      concept __can_flatten_view = requires {
        ranges::flatten_view(std::declval<_Args>()...);
      };
    }  // namespace __detail

    struct _Flatten : std::ranges::range_adaptor_closure<_Flatten> {
      template <class... _Tp>
      requires __detail::__can_flatten_view<_Tp...>
      constexpr auto operator() [[nodiscard]] (_Tp&&... __e) const {
        return ranges::flatten_view(std::forward<_Tp>(__e)...);
      }
      static constexpr bool _S_has_simple_call_op = true;
    };
    struct _ProductN {
      template <class... _Tp>
      constexpr auto operator() [[nodiscard]] (_Tp... __e) const {
        return std::views::cartesian_product(std::views::iota(0, __e)...);
      }
    };

    inline constexpr _Flatten flatten{};
    inline constexpr _ProductN product_n{};
  }  // namespace views
}  // namespace mtd
#line 9 "Test/DataStructure/SternBrocotTree.test.cpp"
// end:tag includes


using ll = long long;

signed main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(0);

  int t;
  std::cin >> t;
  mtd::SternBrocotTree<ll> sbt;
  for (auto _ : std::views::iota(0, t)) {
    std::string s;
    std::cin >> s;

    if (s == "ENCODE_PATH") {
      ll a, b;
      std::cin >> a >> b;
      auto path_rle = sbt.encode(a, b);
      std::cout << path_rle.size() << (path_rle.empty() ? "" : " ");
      for (const auto& [i, right, k] :
           path_rle | std::views::enumerate | mtd::views::flatten) {
        std::cout << (right ? 'R' : 'L') << " " << k
                  << (i == path_rle.size() - 1 ? "" : " ");
      }
      std::cout << std::endl;
    }

    if (s == "DECODE_PATH") {
      ll k;
      std::cin >> k;
      std::vector<std::tuple<bool, ll>> path_rle;
      for (auto __ : std::views::iota(0, k)) {
        char c;
        ll n;
        std::cin >> c >> n;
        path_rle.emplace_back(c == 'R', n);
      }
      auto [a, b] = sbt.decode(path_rle).get();
      std::cout << a << " " << b << std::endl;
    }

    if (s == "LCA") {
      ll a, b, c, d;
      std::cin >> a >> b >> c >> d;
      auto [f, g] = sbt.lca(a, b, c, d).get();
      std::cout << f << " " << g << std::endl;
    }

    if (s == "ANCESTOR") {
      ll k, a, b;
      std::cin >> k >> a >> b;
      try {
        auto [f, g] = sbt.ancestor(a, b, k).get();
        std::cout << f << " " << g << std::endl;
      } catch (const std::runtime_error& e) { std::cout << -1 << std::endl; }
    }

    if (s == "RANGE") {
      ll a, b;
      std::cin >> a >> b;
      try {
        auto [node_l, node_r] = sbt.range(a, b);
        auto [f, g] = node_l.get();
        auto [h, k] = node_r.get();
        std::cout << f << " " << g << " " << h << " " << k << std::endl;
      } catch (const std::runtime_error& e) {
        std::cout << e.what() << std::endl;
      }
    }
  }
}

Back to top page