pell_equation/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::{
4    fmt,
5    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign},
6    sync::Arc,
7};
8type Q = rug::Rational;
9type Z = rug::Integer;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12struct QuadraticField {
13    a: Q,
14    b: Q,
15    d: Arc<Z>,
16}
17impl fmt::Display for QuadraticField {
18    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
19        write!(f, "{}+{}√{}", self.a, self.b, self.d)
20    }
21}
22#[auto_impl_ops::auto_ops]
23impl AddAssign<&QuadraticField> for QuadraticField {
24    fn add_assign(&mut self, other: &Self) {
25        assert_eq!(self.d, other.d);
26        self.a += &other.a;
27        self.b += &other.b;
28    }
29}
30#[auto_impl_ops::auto_ops]
31impl SubAssign<&QuadraticField> for QuadraticField {
32    fn sub_assign(&mut self, other: &Self) {
33        assert_eq!(self.d, other.d);
34        self.a -= &other.a;
35        self.b -= &other.b;
36    }
37}
38#[auto_impl_ops::auto_ops]
39impl SubAssign<&Z> for QuadraticField {
40    fn sub_assign(&mut self, other: &Z) {
41        self.a -= other;
42    }
43}
44#[auto_impl_ops::auto_ops]
45impl MulAssign<&QuadraticField> for QuadraticField {
46    fn mul_assign(&mut self, other: &Self) {
47        assert_eq!(self.d, other.d);
48        let new_a = Q::from(&self.a * &other.a) + Q::from(&self.b * &other.b) * &*self.d;
49        self.b = Q::from(&self.a * &other.b) + Q::from(&self.b * &other.a);
50        self.a = new_a;
51    }
52}
53#[auto_impl_ops::auto_ops]
54impl DivAssign<&QuadraticField> for QuadraticField {
55    fn div_assign(&mut self, other: &Self) {
56        assert_eq!(self.d, other.d);
57        let denom = Q::from(other.a.square_ref()) - Q::from(other.b.square_ref()) * &*self.d;
58        let new_a = (Q::from(&self.a * &other.a) - Q::from(&self.b * &other.b) * &*self.d) / &denom;
59        self.b = (Q::from(&self.b * &other.a) - Q::from(&self.a * &other.b)) / denom;
60        self.a = new_a;
61    }
62}
63impl QuadraticField {
64    fn from_integer(v: Z, d: Arc<Z>) -> Self {
65        Self {
66            a: Q::from(v),
67            b: Q::ZERO.clone(),
68            d,
69        }
70    }
71    fn one(d: Arc<Z>) -> Self {
72        Self::from_integer(Z::ONE.clone(), d)
73    }
74    fn recip(&self) -> Self {
75        QuadraticField::one(Arc::clone(&self.d)) / self
76    }
77    fn floor(&self) -> Z {
78        let d = Z::from(self.d.sqrt_ref());
79        let mut left = &self.a + Q::from(&self.b * &d);
80        let mut right = &self.a + Q::from(&self.b * (d + 1));
81        if right < left {
82            std::mem::swap(&mut left, &mut right);
83        }
84        let offset = Q::from(self.b.square_ref()) * &*self.d;
85        while Q::from(left.floor_ref()) != Q::from(right.floor_ref()) {
86            let mid = Q::from(&left + &right) / 2;
87            let y = Q::from(&mid - &self.a).square() - &offset;
88            if y <= *Q::ZERO {
89                left = mid;
90            } else {
91                right = mid;
92            }
93        }
94        left.floor().numer().clone()
95    }
96}
97
98/// Calculate continued fraction of √d
99///
100/// Calculate [simple continued fraction](https://en.wikipedia.org/wiki/Simple_continued_fraction)
101/// of √d.  
102/// ex : √2 = [1; 2, 2, 2, ...]
103/// ```
104/// use rug::Integer;
105/// let v = pell_equation::continued_fraction_of_sqrt(Integer::from(2));
106/// assert_eq!(v, vec![Integer::from(1), Integer::from(2)]);
107/// ```
108pub fn continued_fraction_of_sqrt(d: Z) -> Vec<Z> {
109    let sd = Z::from(d.sqrt_ref());
110    if Z::from(sd.square_ref()) == d {
111        return vec![sd];
112    }
113    let d = Arc::new(d);
114    let mut v = QuadraticField {
115        a: Q::ZERO.clone(),
116        b: Q::ONE.clone(),
117        d: Arc::clone(&d),
118    };
119    let int = v.floor();
120    v = (v - &int).recip();
121    let v0 = v.clone();
122    let mut a = vec![int];
123    loop {
124        let int = v.floor();
125        v = (v - &int).recip();
126        a.push(int);
127        if v == v0 {
128            return a;
129        }
130    }
131}
132
133/// Fundamental solution of `x^2 - d*y^2 = ±1`
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub enum Solution {
136    /// Fundamental solution of `x^2 - d*y^2 = -1`
137    Negative(Z, Z),
138    /// Fundamental solution of `x^2 - d*y^2 = 1`
139    Positive(Z, Z),
140}
141
142fn solve_pell_aux(a: Vec<Z>, d: Z) -> Solution {
143    let mut p_old = Z::ONE.clone();
144    let mut q_old = Z::ZERO;
145    let mut p_now = a[0].clone();
146    let mut q_now = Z::ONE.clone();
147    // println!("{p_old} {q_old}");
148    // println!("{p_now} {q_now}");
149    let z = Q::from(p_now.square_ref()) - Q::from(q_now.square_ref()) * &d;
150    if z == -Z::ONE.clone() {
151        return Solution::Negative(p_now, q_now);
152    } else if z == *Z::ONE {
153        return Solution::Positive(p_now, q_now);
154    }
155    for a in a.iter().skip(1) {
156        let p_new = a * &p_now + p_old;
157        let q_new = a * &q_now + q_old;
158        p_old = p_now;
159        q_old = q_now;
160        p_now = p_new;
161        q_now = q_new;
162        // println!("{p_now} {q_now}");
163        let z = Q::from(p_now.square_ref()) - Q::from(q_now.square_ref()) * &d;
164        if z == -Z::ONE.clone() {
165            return Solution::Negative(p_now, q_now);
166        } else if z == *Z::ONE {
167            return Solution::Positive(p_now, q_now);
168        }
169    }
170    unreachable!()
171}
172
173/// Calculate fundamental solution of `x^2 - d*y^2 = ±1`
174///
175/// If `x^2 - d*y^2 = -1` has nontrivial solution, returns its fundamental solution.  
176/// Otherwise returns fundamental solution of `x^2 - d*y^2 = 1`.
177/// ```
178/// use rug::Integer;
179/// let v = pell_equation::solve_pell(Integer::from(2));
180/// assert_eq!(v, pell_equation::Solution::Negative(Integer::from(1), Integer::from(1)));
181/// let w = pell_equation::solve_pell(Integer::from(3));
182/// assert_eq!(w, pell_equation::Solution::Positive(Integer::from(2), Integer::from(1)));
183/// ```
184pub fn solve_pell(d: Z) -> Solution {
185    let a = continued_fraction_of_sqrt(d.clone());
186    solve_pell_aux(a, d)
187}
188
189/// Calculate fundamental solution of `x^2 - d*y^2 = -1`
190///
191/// If `x^2 - d*y^2 = -1` has nontrivial solution, returns its fundamental solution.  
192/// Otherwise returns None.
193/// ```
194/// use rug::Integer;
195/// let v = pell_equation::solve_pell_negative(Integer::from(2));
196/// assert_eq!(v, Some((Integer::from(1), Integer::from(1))));
197/// let w = pell_equation::solve_pell_negative(Integer::from(3));
198/// assert_eq!(w, None);
199/// ```
200pub fn solve_pell_negative(d: Z) -> Option<(Z, Z)> {
201    let a = continued_fraction_of_sqrt(d.clone());
202    if (a.len() - 1) % 2 == 0 {
203        return None;
204    }
205    let Solution::Negative(x, y) = solve_pell_aux(a, d) else {
206        unreachable!()
207    };
208    Some((x, y))
209}
210
211/// Calculate fundamental solution of `x^2 - d*y^2 = 1`
212///
213/// This function returns `x^2 - d*y^2 = 1` fundamental solution.  
214/// ```
215/// use rug::Integer;
216/// let v = pell_equation::solve_pell_positive(Integer::from(2));
217/// assert_eq!(v, (Integer::from(3), Integer::from(2)));
218/// let w = pell_equation::solve_pell_positive(Integer::from(3));
219/// assert_eq!(w, (Integer::from(2), Integer::from(1)));
220/// ```
221pub fn solve_pell_positive(d: Z) -> (Z, Z) {
222    match solve_pell(d.clone()) {
223        Solution::Positive(x, y) => (x, y),
224        Solution::Negative(x, y) => {
225            let q = QuadraticField {
226                a: Q::from(x),
227                b: Q::from(y),
228                d: d.into(),
229            };
230            let q2 = &q * &q;
231            (q2.a.numer().clone(), q2.b.numer().clone())
232        }
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    #[test]
241    fn test() {
242        let d = Arc::new(Z::from(3));
243        let q = QuadraticField {
244            a: Q::ZERO.clone(),
245            b: Q::ONE.clone(),
246            d: Arc::clone(&d),
247        };
248        let a = q.floor();
249        assert_eq!(&a, Z::ONE);
250        let q = (q - a).recip();
251        let r = QuadraticField {
252            a: Q::from_f64(0.5).unwrap(),
253            b: Q::from_f64(0.5).unwrap(),
254            d: Arc::clone(&d),
255        };
256        assert_eq!(q, r);
257    }
258    fn to_z(v: &[i32]) -> Vec<Z> {
259        v.iter().map(|x| Z::from(*x)).collect()
260    }
261    // https://planetmath.org/tableofcontinuedfractionsofsqrtnfor1n102
262    #[test]
263    fn test_continued_fraction_of_sqrt2() {
264        let v = continued_fraction_of_sqrt(Z::from(2));
265        assert_eq!(v, to_z(&[1, 2]));
266    }
267    #[test]
268    fn test_continued_fraction_of_sqrt3() {
269        let v = continued_fraction_of_sqrt(Z::from(3));
270        assert_eq!(v, to_z(&[1, 1, 2]));
271    }
272    #[test]
273    fn test_continued_fraction_of_sqrt5() {
274        let v = continued_fraction_of_sqrt(Z::from(5));
275        assert_eq!(v, to_z(&[2, 4]));
276    }
277    #[test]
278    fn test_continued_fraction_of_sqrt6() {
279        let v = continued_fraction_of_sqrt(Z::from(6));
280        assert_eq!(v, to_z(&[2, 2, 4]));
281    }
282    #[test]
283    fn test_continued_fraction_of_sqrt7() {
284        let v = continued_fraction_of_sqrt(Z::from(7));
285        assert_eq!(v, to_z(&[2, 1, 1, 1, 4]));
286    }
287    #[test]
288    fn test_continued_fraction_of_sqrt8() {
289        let v = continued_fraction_of_sqrt(Z::from(8));
290        assert_eq!(v, to_z(&[2, 1, 4]));
291    }
292    #[test]
293    fn test_continued_fraction_of_sqrt10() {
294        let v = continued_fraction_of_sqrt(Z::from(10));
295        assert_eq!(v, to_z(&[3, 6]));
296    }
297    #[test]
298    fn test_continued_fraction_of_sqrt11() {
299        let v = continued_fraction_of_sqrt(Z::from(11));
300        assert_eq!(v, to_z(&[3, 3, 6]));
301    }
302    #[test]
303    fn test_continued_fraction_of_sqrt12() {
304        let v = continued_fraction_of_sqrt(Z::from(12));
305        assert_eq!(v, to_z(&[3, 2, 6]));
306    }
307    #[test]
308    fn test_continued_fraction_of_sqrt13() {
309        let v = continued_fraction_of_sqrt(Z::from(13));
310        assert_eq!(v, to_z(&[3, 1, 1, 1, 1, 6]));
311    }
312    #[test]
313    fn test_continued_fraction_of_sqrt31() {
314        let v = continued_fraction_of_sqrt(Z::from(31));
315        assert_eq!(v, to_z(&[5, 1, 1, 3, 5, 3, 1, 1, 10]));
316    }
317    #[test]
318    fn test_continued_fraction_of_sqrt94() {
319        let v = continued_fraction_of_sqrt(Z::from(94));
320        assert_eq!(
321            v,
322            to_z(&[9, 1, 2, 3, 1, 1, 5, 1, 8, 1, 5, 1, 1, 3, 2, 1, 18])
323        );
324    }
325    #[test]
326    fn test_continued_fraction_of_sqrt338() {
327        let v = continued_fraction_of_sqrt(Z::from(338));
328        assert_eq!(v, to_z(&[18, 2, 1, 1, 2, 36]));
329    }
330    #[test]
331    fn test_solve_pell() {
332        let v = solve_pell(Z::from(653));
333        assert_eq!(
334            v,
335            Solution::Negative(Z::from(2291286382u64), Z::from(89664965))
336        );
337    }
338    #[test]
339    fn test_solve_pell2() {
340        let v = solve_pell(Z::from(115));
341        assert_eq!(v, Solution::Positive(Z::from(1126), Z::from(105)));
342    }
343}
pell_equation/lib.rs

pell_equation/
lib.rs
