competitive/math/
polynomial.rs

1use crate::num::{One, Zero};
2
3#[codesnip::entry("Polynomial", include("zero_one"))]
4#[derive(Clone, Debug, Default, PartialEq, Eq)]
5pub struct Polynomial<T> {
6    pub data: Vec<T>,
7}
8#[codesnip::entry("Polynomial")]
9mod polynomial_impls {
10    use super::*;
11    use std::ops::{Add, Div, Index, IndexMut, Mul, Rem, Sub};
12    impl<T> Polynomial<T> {
13        pub fn from_vec(data: Vec<T>) -> Self {
14            Self { data }
15        }
16        pub fn length(&self) -> usize {
17            self.data.len()
18        }
19    }
20    impl<T> Zero for Polynomial<T> {
21        fn zero() -> Self {
22            Self::from_vec(Vec::new())
23        }
24    }
25    impl<T: Zero + One> One for Polynomial<T> {
26        fn one() -> Self {
27            Self::from_vec(vec![Zero::zero(), One::one()])
28        }
29    }
30    impl<T: Clone + Zero + Add<Output = T> + Mul<Output = T>> Polynomial<T> {
31        pub fn assign(&self, x: T) -> T {
32            let mut res = Zero::zero();
33            for c in self.data.iter().rev().cloned() {
34                res = res * x.clone() + c;
35            }
36            res
37        }
38    }
39    impl<T> Index<usize> for Polynomial<T> {
40        type Output = T;
41        fn index(&self, index: usize) -> &Self::Output {
42            &self.data[index]
43        }
44    }
45    impl<T> IndexMut<usize> for Polynomial<T> {
46        fn index_mut(&mut self, index: usize) -> &mut Self::Output {
47            &mut self.data[index]
48        }
49    }
50    impl<T: Copy + Add<Output = T>> Add<&Polynomial<T>> for &Polynomial<T> {
51        type Output = Polynomial<T>;
52        fn add(self, rhs: &Polynomial<T>) -> Self::Output {
53            let (x, y) = if self.length() < rhs.length() {
54                (rhs, self)
55            } else {
56                (self, rhs)
57            };
58            let mut x = x.clone();
59            for j in 0..y.length() {
60                x[j] = x[j] + y[j];
61            }
62            x
63        }
64    }
65    impl<T: Copy + Sub<Output = T>> Sub<&Polynomial<T>> for &Polynomial<T> {
66        type Output = Polynomial<T>;
67        fn sub(self, rhs: &Polynomial<T>) -> Self::Output {
68            let (x, y) = if self.length() < rhs.length() {
69                (rhs, self)
70            } else {
71                (self, rhs)
72            };
73            let mut x = x.clone();
74            for j in 0..y.length() {
75                x[j] = x[j] - y[j];
76            }
77            x
78        }
79    }
80    impl<T: Copy + Zero + Add<Output = T> + Mul<Output = T>> Mul<&Polynomial<T>> for &Polynomial<T> {
81        type Output = Polynomial<T>;
82        fn mul(self, rhs: &Polynomial<T>) -> Self::Output {
83            let mut res =
84                Polynomial::from_vec(vec![Zero::zero(); self.length() + rhs.length() - 1]);
85            for i in 0..self.length() {
86                for j in 0..rhs.length() {
87                    res[i + j] = res[i + j] + self[i] * rhs[j];
88                }
89            }
90            res
91        }
92    }
93    impl<T: Copy + Zero + Sub<Output = T> + Mul<Output = T> + Div<Output = T>> Div<&Polynomial<T>>
94        for &Polynomial<T>
95    {
96        type Output = Polynomial<T>;
97        fn div(self, rhs: &Polynomial<T>) -> Self::Output {
98            let mut x = self.clone();
99            let mut res = Polynomial::from_vec(vec![]);
100            for i in (rhs.length() - 1..x.length()).rev() {
101                let t = x[i] / rhs[rhs.length() - 1];
102                res.data.push(t);
103                for j in 0..rhs.length() {
104                    x[i - j] = x[i - j] - t * rhs[rhs.length() - 1 - j];
105                }
106            }
107            res.data.reverse();
108            res
109        }
110    }
111    impl<T: Copy + Zero + Sub<Output = T> + Mul<Output = T> + Div<Output = T>> Rem<&Polynomial<T>>
112        for &Polynomial<T>
113    {
114        type Output = Polynomial<T>;
115        fn rem(self, rhs: &Polynomial<T>) -> Self::Output {
116            let mut x = self.clone();
117            for i in (rhs.length() - 1..x.length()).rev() {
118                let t = x[i] / rhs[rhs.length() - 1];
119                for j in 0..rhs.length() {
120                    x[i - j] = x[i - j] - t * rhs[rhs.length() - 1 - j];
121                }
122            }
123            x.data.truncate(rhs.length() - 1);
124            x
125        }
126    }
127    impl<T: Copy + Zero + One + Add<Output = T> + Mul<Output = T>> Polynomial<T> {
128        pub fn pow(&self, mut n: usize) -> Self {
129            let mut x = self.clone();
130            let mut res = Self::one();
131            while n > 0 {
132                if n & 1 == 1 {
133                    res = &res * &x;
134                }
135                x = &x * &x;
136                n >>= 1;
137            }
138            res
139        }
140    }
141}
competitive/math/polynomial.rs

competitive/math/
polynomial.rs
