Skip to main content

sirius/solver/poly/
arithmetic.rs

1use super::{Coef, Mono, Poly};
2use std::cmp::Ordering;
3
4/// Exact monomial division with coefficients.
5pub fn monomial_div(lhs: &(Coef, Mono), rhs: &(Coef, Mono)) -> Option<(Coef, Mono)> {
6    if rhs.0.is_zero() {
7        None
8    } else if lhs.0.is_zero() {
9        Some((Coef::new(0), Mono::unit()))
10    } else if let Some(quot) = lhs.1.div(&rhs.1) {
11        let const_quot = Coef::new(lhs.0.get() / rhs.0.get());
12        if rhs.0 * const_quot == lhs.0 {
13            Some((const_quot, quot))
14        } else {
15            None
16        }
17    } else {
18        None
19    }
20}
21
22impl super::Poly {
23    pub fn add(&self, rhs: &Self) -> Self {
24        let mut terms = Vec::with_capacity(self.terms.len() + rhs.terms.len());
25        let (mut a, mut b) = (self.terms.iter().peekable(), rhs.terms.iter().peekable());
26        loop {
27            match (a.peek(), b.peek()) {
28                (Some(&(ac, am)), Some(&(bc, bm))) => match am.cmp(bm) {
29                    Ordering::Greater => {
30                        terms.push((*ac, am.clone()));
31                        a.next();
32                    }
33                    Ordering::Less => {
34                        terms.push((*bc, bm.clone()));
35                        b.next();
36                    }
37                    Ordering::Equal => {
38                        let c = ac + bc;
39                        if c != 0.into() {
40                            terms.push((c, am.clone()));
41                        }
42                        a.next();
43                        b.next();
44                    }
45                },
46                (Some(_), None) => terms.extend(a.by_ref().cloned()),
47                (None, Some(_)) => terms.extend(b.by_ref().cloned()),
48                (None, None) => break,
49            }
50        }
51        Self { terms }.debug_checked()
52    }
53
54    pub fn mul(&self, rhs: &Self) -> Self {
55        let mut buf = Vec::with_capacity(self.terms.len() * rhs.terms.len());
56        for (ac, am) in &self.terms {
57            for (bc, bm) in &rhs.terms {
58                buf.push((ac * bc, (am.mul(bm))));
59            }
60        }
61        Self::from_terms(buf)
62    }
63
64    pub fn mul_scalar<T: Into<Coef>>(&self, c: T) -> Self {
65        let c = c.into();
66        if c.is_zero() {
67            Self::zero()
68        } else {
69            Self {
70                terms: self.terms.iter().map(|t| (t.0 * c, t.1.clone())).collect(),
71            }
72        }
73    }
74
75    pub fn pow(&self, mut e: u32) -> Self {
76        let mut acc = Self::constant(1);
77        let mut base = self.clone();
78        while e > 0 {
79            if e & 1 == 1 {
80                acc = acc.mul(&base);
81            }
82            e >>= 1;
83            if e > 0 {
84                base = base.mul(&base);
85            }
86        }
87        acc
88    }
89
90    pub fn neg(&self) -> Self {
91        self.mul_scalar(-1)
92    }
93
94    pub fn sub(&self, rhs: &Self) -> Self {
95        self.add(&rhs.neg())
96    }
97
98    /// Via [polynomial long division](https://en.wikipedia.org/wiki/Polynomial_long_division).
99    pub fn compound_divide(&self, divisors: &[Self]) -> (Vec<Self>, Self) {
100        if divisors.is_empty() {
101            return (vec![], self.clone());
102        }
103
104        let mut dividend = self.clone();
105
106        let mut rem = Poly::constant(0);
107        let mut quotients: Vec<Vec<(Coef, Mono)>> =
108            std::iter::repeat_n(Vec::default(), divisors.len()).collect();
109
110        let mut curr_term = 0;
111        let mut curr_divisor = 0;
112
113        while dividend.terms.len() > curr_term {
114            let self_lt = dividend.terms[curr_term].clone();
115            if !divisors[curr_divisor].is_zero() {
116                let div_lt = &divisors[curr_divisor].leading_term();
117                let self_over_div_lt = monomial_div(&self_lt, div_lt);
118                if let Some(self_over_div_lt) = self_over_div_lt {
119                    quotients[curr_divisor].push(self_over_div_lt.clone());
120
121                    let self_over_div_lt = Poly {
122                        terms: vec![self_over_div_lt],
123                    };
124
125                    dividend = dividend.sub(&self_over_div_lt.mul(&divisors[curr_divisor]));
126
127                    curr_divisor = 0;
128                } else {
129                    curr_divisor += 1;
130                }
131            } else {
132                curr_divisor += 1;
133            }
134
135            if curr_divisor == divisors.len() {
136                let self_lt = Self {
137                    terms: vec![self_lt.clone()],
138                };
139                curr_term += 1;
140                rem = rem.add(&self_lt);
141                curr_divisor = 0;
142            }
143        }
144
145        let quotients = quotients
146            .into_iter()
147            .map(|v| Self::debug_checked(Self { terms: v }))
148            .collect();
149
150        (quotients, rem)
151    }
152
153    pub fn try_divide(&self, divisor: &Self) -> Option<Self> {
154        let (mut quots, rem) = self.compound_divide(std::slice::from_ref(divisor));
155
156        if rem.is_zero() {
157            Some(quots.pop().unwrap())
158        } else {
159            None
160        }
161    }
162}