competitive/math/formal_power_series/
formal_power_series_nums.rs

1#![allow(clippy::suspicious_arithmetic_impl, clippy::suspicious_op_assign_impl)]
2
3use super::*;
4use std::{
5    mem::take,
6    ops::{
7        Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Shl, ShlAssign, Shr,
8        ShrAssign, Sub, SubAssign,
9    },
10};
11
12impl<T, C> AddAssign<T> for FormalPowerSeries<T, C>
13where
14    T: FormalPowerSeriesCoefficient,
15{
16    fn add_assign(&mut self, rhs: T) {
17        if self.length() == 0 {
18            self.data.push(T::zero());
19        }
20        self.data[0].add_assign(rhs);
21    }
22}
23impl<T, C> SubAssign<T> for FormalPowerSeries<T, C>
24where
25    T: FormalPowerSeriesCoefficient,
26{
27    fn sub_assign(&mut self, rhs: T) {
28        if self.length() == 0 {
29            self.data.push(T::zero());
30        }
31        self.data[0].sub_assign(rhs);
32        self.trim_tail_zeros();
33    }
34}
35impl<T, C> MulAssign<T> for FormalPowerSeries<T, C>
36where
37    T: FormalPowerSeriesCoefficient,
38{
39    fn mul_assign(&mut self, rhs: T) {
40        for x in self.iter_mut() {
41            x.mul_assign(&rhs);
42        }
43    }
44}
45impl<T, C> DivAssign<T> for FormalPowerSeries<T, C>
46where
47    T: FormalPowerSeriesCoefficient,
48{
49    fn div_assign(&mut self, rhs: T) {
50        let rinv = T::one() / rhs;
51        for x in self.iter_mut() {
52            x.mul_assign(&rinv);
53        }
54    }
55}
56macro_rules! impl_fps_single_binop {
57    ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
58        impl<T, C> $imp_assign<&T> for FormalPowerSeries<T, C>
59        where
60            T: FormalPowerSeriesCoefficient,
61        {
62            fn $method_assign(&mut self, rhs: &T) {
63                $imp_assign::$method_assign(self, rhs.clone());
64            }
65        }
66        impl<T, C> $imp<T> for FormalPowerSeries<T, C>
67        where
68            T: FormalPowerSeriesCoefficient,
69        {
70            type Output = Self;
71            fn $method(mut self, rhs: T) -> Self::Output {
72                $imp_assign::$method_assign(&mut self, rhs);
73                self
74            }
75        }
76        impl<T, C> $imp<&T> for FormalPowerSeries<T, C>
77        where
78            T: FormalPowerSeriesCoefficient,
79        {
80            type Output = Self;
81            fn $method(mut self, rhs: &T) -> Self::Output {
82                $imp_assign::$method_assign(&mut self, rhs);
83                self
84            }
85        }
86        impl<T, C> $imp<T> for &FormalPowerSeries<T, C>
87        where
88            T: FormalPowerSeriesCoefficient,
89        {
90            type Output = FormalPowerSeries<T, C>;
91            fn $method(self, rhs: T) -> Self::Output {
92                $imp::$method(self.clone(), rhs)
93            }
94        }
95        impl<T, C> $imp<&T> for &FormalPowerSeries<T, C>
96        where
97            T: FormalPowerSeriesCoefficient,
98        {
99            type Output = FormalPowerSeries<T, C>;
100            fn $method(self, rhs: &T) -> Self::Output {
101                $imp::$method(self.clone(), rhs)
102            }
103        }
104    };
105}
106impl_fps_single_binop!(Add, add, AddAssign, add_assign);
107impl_fps_single_binop!(Sub, sub, SubAssign, sub_assign);
108impl_fps_single_binop!(Mul, mul, MulAssign, mul_assign);
109impl_fps_single_binop!(Div, div, DivAssign, div_assign);
110
111impl<T, C> AddAssign<&Self> for FormalPowerSeries<T, C>
112where
113    T: FormalPowerSeriesCoefficient,
114{
115    fn add_assign(&mut self, rhs: &Self) {
116        if self.length() < rhs.length() {
117            self.resize(rhs.length());
118        }
119        for (x, y) in self.iter_mut().zip(rhs.iter()) {
120            x.add_assign(y);
121        }
122    }
123}
124impl<T, C> SubAssign<&Self> for FormalPowerSeries<T, C>
125where
126    T: FormalPowerSeriesCoefficient,
127{
128    fn sub_assign(&mut self, rhs: &Self) {
129        if self.length() < rhs.length() {
130            self.resize(rhs.length());
131        }
132        for (x, y) in self.iter_mut().zip(rhs.iter()) {
133            x.sub_assign(y);
134        }
135        self.trim_tail_zeros();
136    }
137}
138
139macro_rules! impl_fps_binop_addsub {
140    ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
141        impl<T, C> $imp_assign for FormalPowerSeries<T, C>
142        where
143            T: FormalPowerSeriesCoefficient,
144        {
145            fn $method_assign(&mut self, rhs: Self) {
146                $imp_assign::$method_assign(self, &rhs);
147            }
148        }
149        impl<T, C> $imp for FormalPowerSeries<T, C>
150        where
151            T: FormalPowerSeriesCoefficient,
152        {
153            type Output = Self;
154            fn $method(mut self, rhs: Self) -> Self::Output {
155                $imp_assign::$method_assign(&mut self, &rhs);
156                self
157            }
158        }
159        impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
160        where
161            T: FormalPowerSeriesCoefficient,
162        {
163            type Output = Self;
164            fn $method(mut self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
165                $imp_assign::$method_assign(&mut self, rhs);
166                self
167            }
168        }
169        impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
170        where
171            T: FormalPowerSeriesCoefficient,
172        {
173            type Output = FormalPowerSeries<T, C>;
174            fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
175                let mut self_ = self.clone();
176                $imp_assign::$method_assign(&mut self_, &rhs);
177                self_
178            }
179        }
180        impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
181        where
182            T: FormalPowerSeriesCoefficient,
183        {
184            type Output = FormalPowerSeries<T, C>;
185            fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
186                let mut self_ = self.clone();
187                $imp_assign::$method_assign(&mut self_, rhs);
188                self_
189            }
190        }
191    };
192}
193impl_fps_binop_addsub!(Add, add, AddAssign, add_assign);
194impl_fps_binop_addsub!(Sub, sub, SubAssign, sub_assign);
195
196impl<T, C> Mul for FormalPowerSeries<T, C>
197where
198    C: ConvolveSteps<T = Vec<T>>,
199{
200    type Output = Self;
201    fn mul(self, rhs: Self) -> Self::Output {
202        Self::from_vec(C::convolve(self.data, rhs.data))
203    }
204}
205impl<T, C> Div for FormalPowerSeries<T, C>
206where
207    T: FormalPowerSeriesCoefficient,
208    C: ConvolveSteps<T = Vec<T>>,
209{
210    type Output = Self;
211    fn div(mut self, mut rhs: Self) -> Self::Output {
212        self.trim_tail_zeros();
213        rhs.trim_tail_zeros();
214        if self.length() < rhs.length() {
215            return Self::zero();
216        }
217        self.data.reverse();
218        rhs.data.reverse();
219        let n = self.length() - rhs.length() + 1;
220        let mut res = self * rhs.inv(n);
221        res.truncate(n);
222        res.data.reverse();
223        res
224    }
225}
226impl<T, C> Rem for FormalPowerSeries<T, C>
227where
228    T: FormalPowerSeriesCoefficient,
229    C: ConvolveSteps<T = Vec<T>>,
230{
231    type Output = Self;
232    fn rem(self, rhs: Self) -> Self::Output {
233        let mut rem = self.clone() - self / rhs.clone() * rhs;
234        rem.trim_tail_zeros();
235        rem
236    }
237}
238
239impl<T, C> FormalPowerSeries<T, C>
240where
241    T: FormalPowerSeriesCoefficient,
242    C: ConvolveSteps<T = Vec<T>>,
243{
244    pub fn div_rem(self, rhs: Self) -> (Self, Self) {
245        let div = self.clone() / rhs.clone();
246        let mut rem = self - div.clone() * rhs;
247        rem.trim_tail_zeros();
248        (div, rem)
249    }
250}
251
252macro_rules! impl_fps_binop_conv {
253    ($imp:ident, $method:ident, $imp_assign:ident, $method_assign:ident) => {
254        impl<T, C> $imp_assign for FormalPowerSeries<T, C>
255        where
256            T: FormalPowerSeriesCoefficient,
257            C: ConvolveSteps<T = Vec<T>>,
258        {
259            fn $method_assign(&mut self, rhs: Self) {
260                *self = $imp::$method(Self::from_vec(take(&mut self.data)), rhs);
261            }
262        }
263        impl<T, C> $imp_assign<&Self> for FormalPowerSeries<T, C>
264        where
265            T: FormalPowerSeriesCoefficient,
266            C: ConvolveSteps<T = Vec<T>>,
267        {
268            fn $method_assign(&mut self, rhs: &Self) {
269                $imp_assign::$method_assign(self, rhs.clone());
270            }
271        }
272        impl<T, C> $imp<&FormalPowerSeries<T, C>> for FormalPowerSeries<T, C>
273        where
274            T: FormalPowerSeriesCoefficient,
275            C: ConvolveSteps<T = Vec<T>>,
276        {
277            type Output = Self;
278            fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
279                $imp::$method(self, rhs.clone())
280            }
281        }
282        impl<T, C> $imp<FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
283        where
284            T: FormalPowerSeriesCoefficient,
285            C: ConvolveSteps<T = Vec<T>>,
286        {
287            type Output = FormalPowerSeries<T, C>;
288            fn $method(self, rhs: FormalPowerSeries<T, C>) -> Self::Output {
289                $imp::$method(self.clone(), rhs)
290            }
291        }
292        impl<T, C> $imp<&FormalPowerSeries<T, C>> for &FormalPowerSeries<T, C>
293        where
294            T: FormalPowerSeriesCoefficient,
295            C: ConvolveSteps<T = Vec<T>>,
296        {
297            type Output = FormalPowerSeries<T, C>;
298            fn $method(self, rhs: &FormalPowerSeries<T, C>) -> Self::Output {
299                $imp::$method(self.clone(), rhs.clone())
300            }
301        }
302    };
303}
304impl_fps_binop_conv!(Mul, mul, MulAssign, mul_assign);
305impl_fps_binop_conv!(Div, div, DivAssign, div_assign);
306impl_fps_binop_conv!(Rem, rem, RemAssign, rem_assign);
307
308impl<T, C> Neg for FormalPowerSeries<T, C>
309where
310    T: FormalPowerSeriesCoefficient,
311{
312    type Output = Self;
313    fn neg(mut self) -> Self::Output {
314        for x in self.iter_mut() {
315            *x = -x.clone();
316        }
317        self
318    }
319}
320impl<T, C> Neg for &FormalPowerSeries<T, C>
321where
322    T: FormalPowerSeriesCoefficient,
323{
324    type Output = FormalPowerSeries<T, C>;
325    fn neg(self) -> Self::Output {
326        self.clone().neg()
327    }
328}
329
330impl<T, C> ShrAssign<usize> for FormalPowerSeries<T, C>
331where
332    T: FormalPowerSeriesCoefficient,
333{
334    fn shr_assign(&mut self, rhs: usize) {
335        if self.length() <= rhs {
336            *self = Self::zero();
337        } else {
338            for i in rhs..self.length() {
339                self[i - rhs] = self[i].clone();
340            }
341            self.truncate(self.length() - rhs);
342        }
343    }
344}
345impl<T, C> ShlAssign<usize> for FormalPowerSeries<T, C>
346where
347    T: FormalPowerSeriesCoefficient,
348{
349    fn shl_assign(&mut self, rhs: usize) {
350        let n = self.length();
351        self.resize(n + rhs);
352        for i in (0..n).rev() {
353            self[i + rhs] = self[i].clone();
354        }
355        for i in 0..rhs {
356            self[i] = T::zero();
357        }
358    }
359}
360
361impl<T, C> Shr<usize> for FormalPowerSeries<T, C>
362where
363    T: FormalPowerSeriesCoefficient,
364{
365    type Output = Self;
366    fn shr(mut self, rhs: usize) -> Self::Output {
367        self.shr_assign(rhs);
368        self
369    }
370}
371impl<T, C> Shl<usize> for FormalPowerSeries<T, C>
372where
373    T: FormalPowerSeriesCoefficient,
374{
375    type Output = Self;
376    fn shl(mut self, rhs: usize) -> Self::Output {
377        self.shl_assign(rhs);
378        self
379    }
380}
381impl<T, C> Shr<usize> for &FormalPowerSeries<T, C>
382where
383    T: FormalPowerSeriesCoefficient,
384{
385    type Output = FormalPowerSeries<T, C>;
386    fn shr(self, rhs: usize) -> Self::Output {
387        if self.length() <= rhs {
388            Self::Output::zero()
389        } else {
390            let mut f = Self::Output::zeros(self.length() - rhs);
391            for i in rhs..self.length() {
392                f[i - rhs] = self[i].clone();
393            }
394            f
395        }
396    }
397}
398impl<T, C> Shl<usize> for &FormalPowerSeries<T, C>
399where
400    T: FormalPowerSeriesCoefficient,
401{
402    type Output = FormalPowerSeries<T, C>;
403    fn shl(self, rhs: usize) -> Self::Output {
404        let mut f = Self::Output::zeros(self.length() + rhs);
405        for (i, x) in self.iter().cloned().enumerate().rev() {
406            f[i + rhs] = x;
407        }
408        f
409    }
410}