competitive/math/
polynomial.rs

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

competitive/math/
polynomial.rs
