Stern–Brocot tree (#322)

Author: hitonanode

number/stern_brocot_tree.hpp

@@ -0,0 +1,97 @@
+#pragma once
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+// Stern–Brocot tree
+// Implementation based on https://miscalc.hatenablog.com/entry/2023/12/22/213007
+namespace SternBrocotTree {
+
+using T = long long;
+
+struct Node {
+    // Subtree contains all rational numbers in (p/q, r/s)
+    T p = 0, q = 1, r = 1, s = 0; // root is (0, \infty)
+
+    // (p + r) / (q + s)
+    T num() const { return p + r; }
+    T den() const { return q + s; }
+};
+
+enum class Direction { Left, Right };
+
+struct Move {
+    Direction dir;
+    T steps;
+};
+
+Node apply(Node node, const Move &mv) {
+    if (mv.dir == Direction::Left) {
+        node.r += node.p * mv.steps;
+        node.s += node.q * mv.steps;
+    } else {
+        node.p += node.r * mv.steps;
+        node.q += node.s * mv.steps;
+    }
+    return node;
+}
+
+// path from root to num/den
+std::vector<Move> encode_path(T num, T den) {
+    std::vector<Move> ret;
+    bool left = false;
+
+    while (num != den) {
+        T steps = num / den;
+        if (den * steps == num) --steps;
+        num -= steps * den;
+        if (steps) ret.push_back({left ? Direction::Left : Direction::Right, steps});
+
+        std::swap(num, den);
+        left = !left;
+    }
+
+    return ret;
+}
+
+Node decode_path(const std::vector<Move> &path) {
+    Node ret{0, 1, 1, 0};
+    for (const Move &mv : path) ret = apply(ret, mv);
+
+    return ret;
+}
+
+std::vector<Move> lca_path(const std::vector<Move> &path1, const std::vector<Move> &path2) {
+    std::vector<Move> ret_path;
+
+    int i1 = 0, i2 = 0;
+    T step1 = path1.empty() ? 0 : path1.front().steps;
+    T step2 = path2.empty() ? 0 : path2.front().steps;
+
+    while (i1 < (int)path1.size() and i2 < (int)path2.size()) {
+        if (!step1) {
+            ++i1;
+            if (i1 < (int)path1.size()) step1 = path1.at(i1).steps;
+        } else if (!step2) {
+            ++i2;
+            if (i2 < (int)path2.size()) step2 = path2.at(i2).steps;
+        } else {
+            if (path1.at(i1).dir != path2.at(i2).dir) break;
+            T steps = std::min(step1, step2);
+            step1 -= steps;
+            step2 -= steps;
+
+            if (ret_path.empty() or ret_path.back().dir != path1.at(i1).dir) {
+                Move move{path1.at(i1).dir, steps};
+                ret_path.push_back(move);
+            } else {
+                ret_path.back().steps += steps;
+            }
+        }
+    }
+
+    return ret_path;
+}
+
+} // namespace SternBrocotTree

number/stern_brocot_tree.md

@@ -0,0 +1,25 @@
+---
+title: Stern–Brocot tree
+documentation_of: ./stern_brocot_tree.hpp
+---
+
+Stern–Brocot tree に関連する処理・探索アルゴリズム．大半の関数は与えられる整数の対数時間で動作する．
+
+## 使用方法
+
+```cpp
+long long a, b;  // coprime
+auto path = SternBrocotTree::encode_path(a, b);  // path from root (= 1) to a/b
+
+auto node = SternBrocotTree::decode_path(path);  // path to rational number
+assert(node.num() == a);
+assert(node.den() == b);
+
+long long c, d;
+auto path2 = SternBrocotTree::encode_path(c, d);
+auto path_lca = SternBrocotTree::lca_path(path, path2);  // path to LCA
+```
+
+## 問題例
+
+- [Library Checker: Stern–Brocot Tree](https://judge.yosupo.jp/problem/stern_brocot_tree)

number/test/stern_brocot_tree.test.cpp

@@ -0,0 +1,84 @@
+#define PROBLEM "https://judge.yosupo.jp/problem/stern_brocot_tree"
+#include "../stern_brocot_tree.hpp"
+
+#include <cassert>
+#include <iostream>
+#include <string>
+#include <vector>
+using namespace std;
+
+int main() {
+    cin.tie(nullptr);
+    ios::sync_with_stdio(false);
+
+    int Q;
+    cin >> Q;
+    while (Q--) {
+        string type;
+        cin >> type;
+        if (type == "ENCODE_PATH") {
+            long long a, b;
+            cin >> a >> b;
+
+            auto path = SternBrocotTree::encode_path(a, b);
+
+            cout << path.size();
+            for (auto mv : path) {
+                cout << ' ' << (mv.dir == SternBrocotTree::Direction::Left ? 'L' : 'R') << ' '
+                     << mv.steps;
+            }
+            cout << '\n';
+        } else if (type == "DECODE_PATH") {
+            std::vector<SternBrocotTree::Move> path;
+            int k;
+            cin >> k;
+            while (k--) {
+                char dir;
+                long long steps;
+                cin >> dir >> steps;
+                path.push_back({dir == 'L' ? SternBrocotTree::Direction::Left
+                                           : SternBrocotTree::Direction::Right,
+                                steps});
+            }
+
+            auto node = SternBrocotTree::decode_path(path);
+            cout << node.num() << ' ' << node.den() << '\n';
+        } else if (type == "LCA") {
+            long long a, b, c, d;
+            cin >> a >> b >> c >> d;
+
+            const auto path1 = SternBrocotTree::encode_path(a, b);
+            const auto path2 = SternBrocotTree::encode_path(c, d);
+            const auto path = SternBrocotTree::lca_path(path1, path2);
+            const auto ret = SternBrocotTree::decode_path(path);
+            cout << ret.num() << ' ' << ret.den() << '\n';
+        } else if (type == "ANCESTOR") {
+            long long k, a, b;
+            cin >> k >> a >> b;
+
+            auto path = SternBrocotTree::encode_path(a, b);
+
+            SternBrocotTree::Node node;
+            for (const auto &mv : path) {
+                long long steps = min(mv.steps, k);
+                k -= steps;
+                node = SternBrocotTree::apply(node, {mv.dir, steps});
+            }
+
+            if (k) {
+                cout << "-1\n";
+            } else {
+                cout << node.num() << ' ' << node.den() << '\n';
+            }
+        } else if (type == "RANGE") {
+            long long a, b;
+            cin >> a >> b;
+
+            auto path = SternBrocotTree::encode_path(a, b);
+            auto node = SternBrocotTree::decode_path(path);
+            cout << node.p << ' ' << node.q << ' ' << node.r << ' ' << node.s << '\n';
+        } else {
+            assert(false);
+        }
+    }
+}