competitive/math/formal_power_series/
formal_power_series_nums.rs

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