haar_lib/math/
stern_brocot.rs

1//! Stern-Brocot木
2//!
3//! # References
4//! - <https://miscalc.hatenablog.com/entry/2023/12/22/213007>
5//!
6//! # Problems
7//! - <https://judge.yosupo.jp/problem/stern_brocot_tree>
8//! - <https://judge.yosupo.jp/problem/rational_approximation>
9
10use crate::math::continued_fraction::*;
11
12/// 分数$\frac{a}{b}$を表す。
13#[derive(Clone, Copy, Debug)]
14pub struct Frac(pub u64, pub u64);
15
16impl PartialEq for Frac {
17    fn eq(&self, Frac(c, d): &Self) -> bool {
18        let Frac(a, b) = self;
19        a * d == b * c
20    }
21}
22impl Eq for Frac {}
23impl PartialOrd for Frac {
24    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
25        let Frac(a, b) = self;
26        let Frac(c, d) = other;
27        (a * d).partial_cmp(&(b * c))
28    }
29}
30
31/// Stern-Brocot木上の頂点を表す。
32#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
33pub struct SBNode(u64, u64, u64, u64);
34
35impl SBNode {
36    /// Stern-Brocot木の根を返す。
37    pub fn root() -> Self {
38        Self(0, 1, 1, 0)
39    }
40
41    /// ノードの内部の4つ組の数を返す。
42    pub fn quadruplet(self) -> (u64, u64, u64, u64) {
43        (self.0, self.1, self.2, self.3)
44    }
45
46    /// 左側の子を`n`回辿った先のノードを返す。
47    pub fn left_child(self, n: u64) -> Self {
48        let Self(p, q, r, s) = self;
49        Self(p, q, p * n + r, q * n + s)
50    }
51
52    /// 右側の子を`n`回辿った先のノードを返す。
53    pub fn right_child(self, n: u64) -> Self {
54        let Self(p, q, r, s) = self;
55        Self(p + r * n, q + s * n, r, s)
56    }
57
58    /// ノードの表す値を[`Frac`]型に変換する。
59    pub fn as_frac(self) -> Frac {
60        let Self(p, q, r, s) = self;
61        Frac(p + r, q + s)
62    }
63
64    /// ノードの表す値を[`f64`]型に変換する。
65    pub fn as_f64(self) -> f64 {
66        let Self(p, q, r, s) = self;
67        ((p + r) as f64) / ((q + s) as f64)
68    }
69
70    /// ノードの子孫ノードの表す値の範囲を返す。
71    pub fn range(self) -> (Frac, Frac) {
72        let Self(p, q, r, s) = self;
73        (Frac(p, q), Frac(r, s))
74    }
75
76    /// [`Frac`]から[`SBNode`]に変換する。
77    pub fn from_frac(a: Frac) -> Option<Self> {
78        Some(SBPath::decode(SBPath::encode(a)?))
79    }
80
81    /// Stern-Brocot木上でのノード`a`とノード`b`のLCAを返す。
82    pub fn lca(a: impl Into<Frac>, b: impl Into<Frac>) -> Option<Self> {
83        let pa = SBPath::encode(a)?;
84        let pb = SBPath::encode(b)?;
85
86        let mut pc = vec![];
87        for (a, b) in std::iter::zip(pa.0, pb.0) {
88            if a == b {
89                pc.push(a);
90            } else {
91                match (a, b) {
92                    (SBMove::L(a), SBMove::L(b)) => pc.push(SBMove::L(a.min(b))),
93                    (SBMove::R(a), SBMove::R(b)) => pc.push(SBMove::R(a.min(b))),
94                    _ => {}
95                }
96                break;
97            }
98        }
99
100        Some(SBPath::decode(SBPath(pc)))
101    }
102
103    /// Stern-Brocot木の根からノード`a`へのパスで、根から深さ`d`のノードを返す。
104    pub fn ancestor(a: impl Into<Frac>, mut d: u64) -> Option<Self> {
105        let path = SBPath::encode(a)?;
106
107        let mut path_d = vec![];
108        for m in path.0 {
109            match m {
110                SBMove::L(n) => {
111                    path_d.push(SBMove::L(n.min(d)));
112                    d = d.saturating_sub(n);
113                }
114                SBMove::R(n) => {
115                    path_d.push(SBMove::R(n.min(d)));
116                    d = d.saturating_sub(n);
117                }
118            }
119            if d == 0 {
120                break;
121            }
122        }
123        if d > 0 {
124            return None;
125        }
126
127        Some(SBPath::decode(SBPath(path_d)))
128    }
129}
130
131impl From<SBNode> for Frac {
132    fn from(value: SBNode) -> Self {
133        value.as_frac()
134    }
135}
136
137/// 子ノードへの移動を表す。
138#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
139pub enum SBMove {
140    /// 左の子を`n`回辿る操作を表す。
141    L(u64),
142    /// 右の子を`n`回辿る操作を表す。
143    R(u64),
144}
145/// Stern-Brocot木の根からノードへのパスを表す。
146#[derive(Clone, Debug, Eq, PartialEq, Hash)]
147pub struct SBPath(pub Vec<SBMove>);
148
149impl SBPath {
150    /// [`Frac`]から[`SBPath`]への変換をする。
151    ///
152    /// `f`の表す分数が、$\frac{n}{0}$または$\frac{0}{n}$であるとき、`None`を返す。
153    pub fn encode(f: impl Into<Frac>) -> Option<Self> {
154        let Frac(p, q) = f.into();
155        if p == 0 || q == 0 {
156            return None;
157        }
158
159        let mut cfe = continued_fraction(p, q)?;
160        *cfe.last_mut().unwrap() -= 1;
161
162        let mut ret = vec![];
163        if let Some(&m) = cfe.first() {
164            if m != 0 {
165                ret.push(SBMove::R(m));
166            }
167        }
168
169        for (i, m) in cfe.into_iter().skip(1).enumerate() {
170            if i % 2 == 0 {
171                ret.push(SBMove::L(m));
172            } else {
173                ret.push(SBMove::R(m));
174            }
175        }
176
177        Some(Self(ret))
178    }
179
180    /// [`SBPath`]から[`SBNode`]への変換をする。
181    pub fn decode(self) -> SBNode {
182        let mut ret = SBNode::root();
183        for m in self.0 {
184            match m {
185                SBMove::L(n) => ret = ret.left_child(n),
186                SBMove::R(n) => ret = ret.right_child(n),
187            }
188        }
189        ret
190    }
191}
haar_lib/math/stern_brocot.rs

haar_lib/math/
stern_brocot.rs
