pell_equation/
lib.rs

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

pell_equation/
lib.rs
