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        debug_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        debug_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        debug_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        debug_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(mut a: Vec<Z>, d: Z) -> Solution {
143    let n = a.len() - 1;
144    a.reverse();
145    let mut p_old = Z::ONE.clone();
146    let mut q_old = Z::ZERO;
147    let mut p_now = a.pop().unwrap();
148    let mut q_now = Z::ONE.clone();
149    // println!("{p_old} {q_old}");
150    // println!("{p_now} {q_now}");
151    for _ in 1..n {
152        let ai = a.pop().unwrap();
153        p_old += &ai * &p_now;
154        q_old += ai * &q_now;
155        std::mem::swap(&mut p_old, &mut p_now);
156        std::mem::swap(&mut q_old, &mut q_now);
157        // println!("{p_now} {q_now}");
158    }
159    if n % 2 == 0 {
160        debug_assert_eq!(
161            Z::from(p_now.square_ref()) - Z::from(q_now.square_ref()) * &d,
162            *Z::ONE
163        );
164        Solution::Positive(p_now, q_now)
165    } else {
166        debug_assert_eq!(
167            Z::from(p_now.square_ref()) - Z::from(q_now.square_ref()) * &d,
168            -Z::ONE.clone()
169        );
170        Solution::Negative(p_now, q_now)
171    }
172}
173
174/// Calculate fundamental solution of `x^2 - d*y^2 = ±1`
175///
176/// If `x^2 - d*y^2 = -1` has nontrivial solution, returns its fundamental solution.  
177/// Otherwise returns fundamental solution of `x^2 - d*y^2 = 1`.
178/// ```
179/// use rug::Integer;
180/// let v = pell_equation::solve_pell(Integer::from(2));
181/// assert_eq!(v, pell_equation::Solution::Negative(Integer::from(1), Integer::from(1)));
182/// let w = pell_equation::solve_pell(Integer::from(3));
183/// assert_eq!(w, pell_equation::Solution::Positive(Integer::from(2), Integer::from(1)));
184/// ```
185pub fn solve_pell(d: Z) -> Solution {
186    let a = continued_fraction_of_sqrt(d.clone());
187    solve_pell_aux(a, d)
188}
189
190/// Calculate fundamental solution of `x^2 - d*y^2 = -1`
191///
192/// If `x^2 - d*y^2 = -1` has nontrivial solution, returns its fundamental solution.  
193/// Otherwise returns None.
194/// ```
195/// use rug::Integer;
196/// let v = pell_equation::solve_pell_negative(Integer::from(2));
197/// assert_eq!(v, Some((Integer::from(1), Integer::from(1))));
198/// let w = pell_equation::solve_pell_negative(Integer::from(3));
199/// assert_eq!(w, None);
200/// ```
201pub fn solve_pell_negative(d: Z) -> Option<(Z, Z)> {
202    let a = continued_fraction_of_sqrt(d.clone());
203    if (a.len() - 1) % 2 == 0 {
204        return None;
205    }
206    let Solution::Negative(x, y) = solve_pell_aux(a, d) else {
207        unreachable!()
208    };
209    Some((x, y))
210}
211
212/// Calculate fundamental solution of `x^2 - d*y^2 = 1`
213///
214/// This function returns `x^2 - d*y^2 = 1` fundamental solution.  
215/// ```
216/// use rug::Integer;
217/// let v = pell_equation::solve_pell_positive(Integer::from(2));
218/// assert_eq!(v, (Integer::from(3), Integer::from(2)));
219/// let w = pell_equation::solve_pell_positive(Integer::from(3));
220/// assert_eq!(w, (Integer::from(2), Integer::from(1)));
221/// ```
222pub fn solve_pell_positive(d: Z) -> (Z, Z) {
223    match solve_pell(d.clone()) {
224        Solution::Positive(x, y) => (x, y),
225        Solution::Negative(x, y) => {
226            let q = QuadraticField {
227                a: Q::from(x),
228                b: Q::from(y),
229                d: d.into(),
230            };
231            let q2 = &q * &q;
232            (q2.a.numer().clone(), q2.b.numer().clone())
233        }
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn test() {
243        let d = Arc::new(Z::from(3));
244        let q = QuadraticField {
245            a: Q::ZERO.clone(),
246            b: Q::ONE.clone(),
247            d: Arc::clone(&d),
248        };
249        let a = q.floor();
250        assert_eq!(&a, Z::ONE);
251        let q = (q - a).recip();
252        let r = QuadraticField {
253            a: Q::from_f64(0.5).unwrap(),
254            b: Q::from_f64(0.5).unwrap(),
255            d: Arc::clone(&d),
256        };
257        assert_eq!(q, r);
258    }
259    fn to_z(v: &[i32]) -> Vec<Z> {
260        v.iter().map(|x| Z::from(*x)).collect()
261    }
262    // https://planetmath.org/tableofcontinuedfractionsofsqrtnfor1n102
263    #[test]
264    fn test_continued_fraction_of_sqrt2() {
265        let v = continued_fraction_of_sqrt(Z::from(2));
266        assert_eq!(v, to_z(&[1, 2]));
267    }
268    #[test]
269    fn test_continued_fraction_of_sqrt3() {
270        let v = continued_fraction_of_sqrt(Z::from(3));
271        assert_eq!(v, to_z(&[1, 1, 2]));
272    }
273    #[test]
274    fn test_continued_fraction_of_sqrt5() {
275        let v = continued_fraction_of_sqrt(Z::from(5));
276        assert_eq!(v, to_z(&[2, 4]));
277    }
278    #[test]
279    fn test_continued_fraction_of_sqrt6() {
280        let v = continued_fraction_of_sqrt(Z::from(6));
281        assert_eq!(v, to_z(&[2, 2, 4]));
282    }
283    #[test]
284    fn test_continued_fraction_of_sqrt7() {
285        let v = continued_fraction_of_sqrt(Z::from(7));
286        assert_eq!(v, to_z(&[2, 1, 1, 1, 4]));
287    }
288    #[test]
289    fn test_continued_fraction_of_sqrt8() {
290        let v = continued_fraction_of_sqrt(Z::from(8));
291        assert_eq!(v, to_z(&[2, 1, 4]));
292    }
293    #[test]
294    fn test_continued_fraction_of_sqrt10() {
295        let v = continued_fraction_of_sqrt(Z::from(10));
296        assert_eq!(v, to_z(&[3, 6]));
297    }
298    #[test]
299    fn test_continued_fraction_of_sqrt11() {
300        let v = continued_fraction_of_sqrt(Z::from(11));
301        assert_eq!(v, to_z(&[3, 3, 6]));
302    }
303    #[test]
304    fn test_continued_fraction_of_sqrt12() {
305        let v = continued_fraction_of_sqrt(Z::from(12));
306        assert_eq!(v, to_z(&[3, 2, 6]));
307    }
308    #[test]
309    fn test_continued_fraction_of_sqrt13() {
310        let v = continued_fraction_of_sqrt(Z::from(13));
311        assert_eq!(v, to_z(&[3, 1, 1, 1, 1, 6]));
312    }
313    #[test]
314    fn test_continued_fraction_of_sqrt31() {
315        let v = continued_fraction_of_sqrt(Z::from(31));
316        assert_eq!(v, to_z(&[5, 1, 1, 3, 5, 3, 1, 1, 10]));
317    }
318    #[test]
319    fn test_continued_fraction_of_sqrt94() {
320        let v = continued_fraction_of_sqrt(Z::from(94));
321        assert_eq!(
322            v,
323            to_z(&[9, 1, 2, 3, 1, 1, 5, 1, 8, 1, 5, 1, 1, 3, 2, 1, 18])
324        );
325    }
326    #[test]
327    fn test_continued_fraction_of_sqrt338() {
328        let v = continued_fraction_of_sqrt(Z::from(338));
329        assert_eq!(v, to_z(&[18, 2, 1, 1, 2, 36]));
330    }
331    #[test]
332    fn test_solve_pell() {
333        let v = solve_pell(Z::from(653));
334        assert_eq!(
335            v,
336            Solution::Negative(Z::from(2291286382u64), Z::from(89664965))
337        );
338    }
339    #[test]
340    fn test_solve_pell2() {
341        let v = solve_pell(Z::from(115));
342        assert_eq!(v, Solution::Positive(Z::from(1126), Z::from(105)));
343    }
344    #[test]
345    fn test_solve_pell3() {
346        let v = solve_pell(Z::from(114));
347        assert_eq!(v, Solution::Positive(Z::from(1025), Z::from(96)));
348    }
349    #[test]
350    fn test_solve_pell4() {
351        let v = solve_pell(Z::from(641));
352        assert_eq!(
353            v,
354            Solution::Negative(Z::from(36120833468u64), Z::from(1426687145))
355        );
356    }
357}
pell_equation/lib.rs

pell_equation/
lib.rs
