competitive/math/formal_power_series/
formal_power_series_impls.rs

1use super::*;
2use std::{
3    cmp::Reverse,
4    collections::BinaryHeap,
5    iter::repeat_with,
6    iter::{FromIterator, once},
7    marker::PhantomData,
8    ops::{Index, IndexMut},
9    slice::{Iter, IterMut},
10};
11
12impl<T, C> FormalPowerSeries<T, C> {
13    pub fn from_vec(data: Vec<T>) -> Self {
14        Self {
15            data,
16            _marker: PhantomData,
17        }
18    }
19    pub fn length(&self) -> usize {
20        self.data.len()
21    }
22    pub fn truncate(&mut self, deg: usize) {
23        self.data.truncate(deg)
24    }
25    pub fn iter(&self) -> Iter<'_, T> {
26        self.data.iter()
27    }
28    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
29        self.data.iter_mut()
30    }
31}
32
33impl<T, C> Clone for FormalPowerSeries<T, C>
34where
35    T: Clone,
36{
37    fn clone(&self) -> Self {
38        Self::from_vec(self.data.clone())
39    }
40}
41impl<T, C> PartialEq for FormalPowerSeries<T, C>
42where
43    T: PartialEq,
44{
45    fn eq(&self, other: &Self) -> bool {
46        self.data.eq(&other.data)
47    }
48}
49impl<T, C> Eq for FormalPowerSeries<T, C> where T: PartialEq {}
50
51impl<T, C> FormalPowerSeries<T, C>
52where
53    T: Zero,
54{
55    pub fn zeros(deg: usize) -> Self {
56        repeat_with(T::zero).take(deg).collect()
57    }
58    pub fn resize(&mut self, deg: usize) {
59        self.data.resize_with(deg, Zero::zero)
60    }
61    pub fn resized(mut self, deg: usize) -> Self {
62        self.resize(deg);
63        self
64    }
65    pub fn reversed(mut self) -> Self {
66        self.data.reverse();
67        self
68    }
69}
70
71impl<T, C> FormalPowerSeries<T, C>
72where
73    T: Zero + Clone,
74{
75    pub fn coeff(&self, deg: usize) -> T {
76        self.data.get(deg).cloned().unwrap_or_else(T::zero)
77    }
78}
79
80impl<T, C> FormalPowerSeries<T, C>
81where
82    T: Zero + PartialEq,
83{
84    pub fn trim_tail_zeros(&mut self) {
85        let mut len = self.length();
86        while len > 0 {
87            if self.data[len - 1].is_zero() {
88                len -= 1;
89            } else {
90                break;
91            }
92        }
93        self.truncate(len);
94    }
95}
96
97impl<T, C> Zero for FormalPowerSeries<T, C>
98where
99    T: PartialEq,
100{
101    fn zero() -> Self {
102        Self::from_vec(Vec::new())
103    }
104}
105impl<T, C> One for FormalPowerSeries<T, C>
106where
107    T: PartialEq + One,
108{
109    fn one() -> Self {
110        Self::from(T::one())
111    }
112}
113
114impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
115    type Item = T;
116    type IntoIter = std::vec::IntoIter<T>;
117    fn into_iter(self) -> Self::IntoIter {
118        self.data.into_iter()
119    }
120}
121impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
122    type Item = &'a T;
123    type IntoIter = Iter<'a, T>;
124    fn into_iter(self) -> Self::IntoIter {
125        self.data.iter()
126    }
127}
128impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
129    type Item = &'a mut T;
130    type IntoIter = IterMut<'a, T>;
131    fn into_iter(self) -> Self::IntoIter {
132        self.data.iter_mut()
133    }
134}
135
136impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
137    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
138        Self::from_vec(iter.into_iter().collect())
139    }
140}
141
142impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
143    type Output = T;
144    fn index(&self, index: usize) -> &Self::Output {
145        &self.data[index]
146    }
147}
148impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
149    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
150        &mut self.data[index]
151    }
152}
153
154impl<T, C> From<T> for FormalPowerSeries<T, C> {
155    fn from(x: T) -> Self {
156        once(x).collect()
157    }
158}
159impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
160    fn from(data: Vec<T>) -> Self {
161        Self::from_vec(data)
162    }
163}
164
165impl<T, C> FormalPowerSeries<T, C>
166where
167    T: FormalPowerSeriesCoefficient,
168{
169    pub fn prefix_ref(&self, deg: usize) -> Self {
170        if deg < self.length() {
171            Self::from_vec(self.data[..deg].to_vec())
172        } else {
173            self.clone()
174        }
175    }
176    pub fn prefix(mut self, deg: usize) -> Self {
177        self.data.truncate(deg);
178        self
179    }
180    pub fn even(mut self) -> Self {
181        let mut keep = false;
182        self.data.retain(|_| {
183            keep = !keep;
184            keep
185        });
186        self
187    }
188    pub fn odd(mut self) -> Self {
189        let mut keep = true;
190        self.data.retain(|_| {
191            keep = !keep;
192            keep
193        });
194        self
195    }
196    pub fn diff(mut self) -> Self {
197        let mut c = T::one();
198        for x in self.iter_mut().skip(1) {
199            *x *= &c;
200            c += T::one();
201        }
202        if self.length() > 0 {
203            self.data.remove(0);
204        }
205        self
206    }
207    pub fn integral(mut self) -> Self {
208        let n = self.length();
209        self.data.insert(0, Zero::zero());
210        let mut fact = Vec::with_capacity(n + 1);
211        let mut c = T::one();
212        fact.push(c.clone());
213        for _ in 1..n {
214            fact.push(fact.last().cloned().unwrap() * c.clone());
215            c += T::one();
216        }
217        let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
218        for x in self.iter_mut().skip(1).rev() {
219            *x *= invf.clone() * fact.pop().unwrap();
220            invf *= c.clone();
221            c -= T::one();
222        }
223        self
224    }
225    pub fn parity_inversion(mut self) -> Self {
226        self.iter_mut()
227            .skip(1)
228            .step_by(2)
229            .for_each(|x| *x = -x.clone());
230        self
231    }
232    pub fn eval(&self, x: T) -> T {
233        let mut base = T::one();
234        let mut res = T::zero();
235        for a in self.iter() {
236            res += base.clone() * a.clone();
237            base *= x.clone();
238        }
239        res
240    }
241}
242
243impl<T, C> FormalPowerSeries<T, C>
244where
245    T: FormalPowerSeriesCoefficient,
246    C: ConvolveSteps<T = Vec<T>>,
247{
248    pub fn inv(&self, deg: usize) -> Self {
249        debug_assert!(!self[0].is_zero());
250        if self.data.iter().filter(|x| !x.is_zero()).count()
251            <= deg.next_power_of_two().trailing_zeros() as usize * 6
252        {
253            let pos: Vec<_> = self
254                .data
255                .iter()
256                .enumerate()
257                .skip(1)
258                .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
259                .collect();
260            let mut f = Self::zeros(deg);
261            f[0] = T::one() / self[0].clone();
262            for i in 1..deg {
263                let mut tot = T::zero();
264                for &j in &pos {
265                    if j > i {
266                        break;
267                    }
268                    tot += self[j].clone() * &f[i - j];
269                }
270                f[i] = -tot * &f[0];
271            }
272            return f;
273        }
274        let mut f = Self::from(T::one() / self[0].clone());
275        let mut i = 1;
276        while i < deg {
277            let g = self.prefix_ref((i * 2).min(deg));
278            let h = f.clone();
279            let mut g = C::transform(g.data, 2 * i);
280            let h = C::transform(h.data, 2 * i);
281            C::multiply(&mut g, &h);
282            let mut g = Self::from_vec(C::inverse_transform(g, 2 * i));
283            g >>= i;
284            let mut g = C::transform(g.data, 2 * i);
285            C::multiply(&mut g, &h);
286            let g = Self::from_vec(C::inverse_transform(g, 2 * i));
287            f.data.extend((-g).into_iter().take(i));
288            i *= 2;
289        }
290        f.truncate(deg);
291        f
292    }
293    pub fn exp(&self, deg: usize) -> Self {
294        debug_assert!(self[0].is_zero());
295        if self.data.iter().filter(|x| !x.is_zero()).count()
296            <= deg.next_power_of_two().trailing_zeros() as usize * 16
297        {
298            let diff = self.clone().diff();
299            let pos: Vec<_> = diff
300                .data
301                .iter()
302                .enumerate()
303                .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
304                .collect();
305            let mf = T::memorized_factorial(deg);
306            let mut f = Self::zeros(deg);
307            f[0] = T::one();
308            for i in 1..deg {
309                let mut tot = T::zero();
310                for &j in &pos {
311                    if j > i - 1 {
312                        break;
313                    }
314                    tot += f[i - 1 - j].clone() * &diff[j];
315                }
316                f[i] = tot * T::memorized_inv(&mf, i);
317            }
318            return f;
319        }
320        let mut f = Self::one();
321        let mut i = 1;
322        while i < deg {
323            let mut g = -f.log(i * 2);
324            g[0] += T::one();
325            for (g, x) in g.iter_mut().zip(self.iter().take(i * 2)) {
326                *g += x.clone();
327            }
328            f = (f * g).prefix(i * 2);
329            i *= 2;
330        }
331        f.prefix(deg)
332    }
333    pub fn log(&self, deg: usize) -> Self {
334        (self.inv(deg) * self.clone().diff()).integral().prefix(deg)
335    }
336    pub fn pow(&self, rhs: usize, deg: usize) -> Self {
337        if rhs == 0 {
338            return Self::from_vec(
339                once(T::one())
340                    .chain(repeat_with(T::zero))
341                    .take(deg)
342                    .collect(),
343            );
344        }
345        if let Some(k) = self.iter().position(|x| !x.is_zero()) {
346            if k >= deg.div_ceil(rhs) {
347                Self::zeros(deg)
348            } else {
349                let deg = deg - k * rhs;
350                let x0 = self[k].clone();
351                let mut f = (self >> k) / &x0;
352                if f.data.iter().filter(|x| !x.is_zero()).count()
353                    <= deg.next_power_of_two().trailing_zeros() as usize * 12
354                {
355                    f = f.pow_sparse1(T::from(rhs), deg);
356                } else {
357                    f = (f.log(deg) * &T::from(rhs)).exp(deg);
358                }
359                f *= x0.pow(rhs);
360                f <<= k * rhs;
361                f
362            }
363        } else {
364            Self::zeros(deg)
365        }
366    }
367    fn pow_sparse1(&self, rhs: T, deg: usize) -> Self {
368        debug_assert!(!self[0].is_zero());
369        let pos: Vec<_> = self
370            .data
371            .iter()
372            .enumerate()
373            .skip(1)
374            .filter_map(|(i, x)| if x.is_zero() { None } else { Some(i) })
375            .collect();
376        let mf = T::memorized_factorial(deg);
377        let mut f = Self::zeros(deg);
378        f[0] = T::one();
379        for i in 1..deg {
380            let mut tot = T::zero();
381            for &j in &pos {
382                if j > i {
383                    break;
384                }
385                tot += (T::from(j) * &rhs - T::from(i - j)) * &self[j] * &f[i - j];
386            }
387            f[i] = tot * T::memorized_inv(&mf, i);
388        }
389        f
390    }
391}
392
393impl<T, C> FormalPowerSeries<T, C>
394where
395    T: FormalPowerSeriesCoefficientSqrt,
396    C: ConvolveSteps<T = Vec<T>>,
397{
398    pub fn sqrt(&self, deg: usize) -> Option<Self> {
399        if self[0].is_zero() {
400            if let Some(k) = self.iter().position(|x| !x.is_zero()) {
401                if k % 2 != 0 {
402                    return None;
403                } else if deg > k / 2 {
404                    return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
405                }
406            }
407        } else {
408            let s = self[0].sqrt_coefficient()?;
409            if self.data.iter().filter(|x| !x.is_zero()).count()
410                <= deg.next_power_of_two().trailing_zeros() as usize * 4
411            {
412                let t = self[0].clone();
413                let mut f = self / t;
414                f = f.pow_sparse1(T::from(1) / T::from(2), deg);
415                f *= s;
416                return Some(f);
417            }
418
419            let mut f = Self::from(s);
420            let inv2 = T::one() / (T::one() + T::one());
421            let mut i = 1;
422            while i < deg {
423                f = (&f + &(self.prefix_ref(i * 2) * f.inv(i * 2))).prefix(i * 2) * &inv2;
424                i *= 2;
425            }
426            f.truncate(deg);
427            return Some(f);
428        }
429        Some(Self::zeros(deg))
430    }
431}
432
433impl<T, C> FormalPowerSeries<T, C>
434where
435    T: FormalPowerSeriesCoefficient,
436    C: ConvolveSteps<T = Vec<T>>,
437{
438    pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
439    where
440        F: FnMut(usize) -> T,
441    {
442        let n = self.length();
443        let mut f = Self::zeros(n);
444        for i in 1..n {
445            if !self[i].is_zero() {
446                for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
447                    if j & 1 != 0 {
448                        f[d] += self[i].clone() * &inverse(j);
449                    } else {
450                        f[d] -= self[i].clone() * &inverse(j);
451                    }
452                }
453            }
454        }
455        f.exp(deg)
456    }
457    pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
458    where
459        F: FnMut(usize) -> T,
460    {
461        let n = self.length();
462        let mut f = Self::zeros(n);
463        for i in 1..n {
464            if !self[i].is_zero() {
465                for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
466                    f[d] += self[i].clone() * &inverse(j);
467                }
468            }
469        }
470        f.exp(deg)
471    }
472    /// [x^n] P(x) / Q(x)
473    pub fn bostan_mori(mut self, mut rhs: Self, mut n: usize) -> T
474    where
475        C: NttReuse<T = Vec<T>>,
476    {
477        let mut res = T::zero();
478        rhs.trim_tail_zeros();
479        if self.length() >= rhs.length() {
480            let r = &self / &rhs;
481            if n < r.length() {
482                res = r[n].clone();
483            }
484            self -= r * &rhs;
485            self.trim_tail_zeros();
486        }
487        let k = rhs.length().next_power_of_two();
488        let mut p = C::transform(self.data, k * 2);
489        let mut q = C::transform(rhs.data, k * 2);
490        while n > 0 {
491            let t = C::even_mul_normal_neg(&q, &q);
492            p = if n.is_multiple_of(2) {
493                C::even_mul_normal_neg(&p, &q)
494            } else {
495                C::odd_mul_normal_neg(&p, &q)
496            };
497            q = t;
498            n /= 2;
499            if n != 0 {
500                if C::MULTIPLE {
501                    p = C::transform(C::inverse_transform(p, k), k * 2);
502                    q = C::transform(C::inverse_transform(q, k), k * 2);
503                } else {
504                    p = C::ntt_doubling(p);
505                    q = C::ntt_doubling(q);
506                }
507            }
508        }
509        let p = C::inverse_transform(p, k);
510        let q = C::inverse_transform(q, k);
511        res + p[0].clone() / q[0].clone()
512    }
513    /// return F(x) where [x^n] P(x) / Q(x) = [x^d-1] P(x) F(x)
514    pub fn bostan_mori_msb(self, n: usize) -> Self {
515        let d = self.length() - 1;
516        if n == 0 {
517            return (Self::one() << (d - 1)) / self[0].clone();
518        }
519        let q = self;
520        let mq = q.clone().parity_inversion();
521        let w = (q * &mq).even().bostan_mori_msb(n / 2);
522        let mut s = Self::zeros(w.length() * 2 - (n % 2));
523        for (i, x) in w.iter().enumerate() {
524            s[i * 2 + (1 - n % 2)] = x.clone();
525        }
526        let len = 2 * d + 1;
527        let ts = C::transform(s.prefix(len).data, len);
528        mq.reversed().middle_product(&ts, len).prefix(d + 1)
529    }
530    /// x^n mod self
531    pub fn pow_mod(self, n: usize) -> Self {
532        let d = self.length() - 1;
533        let q = self.reversed();
534        let u = q.clone().bostan_mori_msb(n);
535        let mut f = (u * q).prefix(d).reversed();
536        f.trim_tail_zeros();
537        f
538    }
539    fn middle_product(self, other: &C::F, deg: usize) -> Self {
540        let n = self.length();
541        let mut s = C::transform(self.reversed().data, deg);
542        C::multiply(&mut s, other);
543        Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
544    }
545    pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T> {
546        let n = points.len();
547        if n <= 32 {
548            return points.iter().map(|p| self.eval(p.clone())).collect();
549        }
550        let mut subproduct_tree = Vec::with_capacity(n * 2);
551        subproduct_tree.resize_with(n, Zero::zero);
552        for x in points {
553            subproduct_tree.push(Self::from_vec(vec![-x.clone(), T::one()]));
554        }
555        for i in (1..n).rev() {
556            subproduct_tree[i] = &subproduct_tree[i * 2] * &subproduct_tree[i * 2 + 1];
557        }
558        let mut uptree_t = Vec::with_capacity(n * 2);
559        uptree_t.resize_with(1, Zero::zero);
560        subproduct_tree.reverse();
561        subproduct_tree.pop();
562        let m = self.length();
563        let v = subproduct_tree.pop().unwrap().reversed().resized(m);
564        let s = C::transform(self.data, m * 2);
565        uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(n).reversed());
566        for i in 1..n {
567            let subl = subproduct_tree.pop().unwrap();
568            let subr = subproduct_tree.pop().unwrap();
569            let (dl, dr) = (subl.length(), subr.length());
570            let len = dl.max(dr) + uptree_t[i].length();
571            let s = C::transform(uptree_t[i].data.to_vec(), len);
572            uptree_t.push(subr.middle_product(&s, len).prefix(dl));
573            uptree_t.push(subl.middle_product(&s, len).prefix(dr));
574        }
575        uptree_t[n..]
576            .iter()
577            .map(|u| u.data.first().cloned().unwrap_or_else(Zero::zero))
578            .collect()
579    }
580    pub fn product_all<I>(iter: I, deg: usize) -> Self
581    where
582        I: IntoIterator<Item = Self>,
583    {
584        let mut heap: BinaryHeap<_> = iter
585            .into_iter()
586            .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
587            .collect();
588        while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
589            if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
590                let z = (x * y).prefix(deg);
591                heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
592            } else {
593                return x;
594            }
595        }
596        Self::one()
597    }
598    pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
599    where
600        I: IntoIterator<Item = (Self, Self)>,
601    {
602        let mut heap: BinaryHeap<_> = iter
603            .into_iter()
604            .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
605            .collect();
606        while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
607            if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
608                let zb = (&xb * &yb).prefix(deg);
609                let za = (xa * yb + ya * xb).prefix(deg);
610                heap.push(PartialIgnoredOrd(
611                    Reverse(za.length().max(zb.length())),
612                    (za, zb),
613                ));
614            } else {
615                return (xa, xb);
616            }
617        }
618        (Self::zero(), Self::one())
619    }
620    pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T
621    where
622        C: NttReuse<T = Vec<T>>,
623    {
624        if let Some(x) = a.get(k) {
625            return x.clone();
626        }
627        let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
628        p.bostan_mori(self, k)
629    }
630    pub fn kth_term(a: Vec<T>, k: usize) -> T
631    where
632        C: NttReuse<T = Vec<T>>,
633    {
634        if let Some(x) = a.get(k) {
635            return x.clone();
636        }
637        Self::from_vec(berlekamp_massey(&a)).kth_term_of_linearly_recurrence(a, k)
638    }
639    /// sum_i a_i exp(b_i x)
640    pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
641    where
642        I: IntoIterator<Item = (T, T)>,
643        F: FnMut(usize) -> T,
644    {
645        let (p, q) = Self::sum_all_rational(
646            iter.into_iter()
647                .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
648            deg,
649        );
650        let mut f = (p * q.inv(deg)).prefix(deg);
651        for i in 0..f.length() {
652            f[i] *= inv_fact(i);
653        }
654        f
655    }
656    /// sum_i (a_i x)^j
657    pub fn sum_of_powers<I>(iter: I, deg: usize) -> Self
658    where
659        I: IntoIterator<Item = T>,
660    {
661        let mut n = T::zero();
662        let prod = Self::product_all(
663            iter.into_iter().map(|a| {
664                n += T::one();
665                Self::from_vec(vec![T::one(), -a])
666            }),
667            deg,
668        );
669        (-prod.log(deg).diff() << 1) + Self::from_vec(vec![n])
670    }
671}
672
673impl<M, C> FormalPowerSeries<MInt<M>, C>
674where
675    M: MIntConvert<usize>,
676    C: ConvolveSteps<T = Vec<MInt<M>>>,
677{
678    /// f(x) <- f(x + a)
679    pub fn taylor_shift(mut self, a: MInt<M>, f: &MemorizedFactorial<M>) -> Self {
680        let n = self.length();
681        for i in 0..n {
682            self.data[i] *= f.fact[i];
683        }
684        self.data.reverse();
685        let mut b = a;
686        let mut g = Self::from_vec(f.inv_fact[..n].to_vec());
687        for i in 1..n {
688            g[i] *= b;
689            b *= a;
690        }
691        self *= g;
692        self.truncate(n);
693        self.data.reverse();
694        for i in 0..n {
695            self.data[i] *= f.inv_fact[i];
696        }
697        self
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use super::*;
704    use crate::{num::mint_basic::Modulo1000000009, rand, tools::Xorshift};
705
706    #[test]
707    fn test_bostan_mori() {
708        let mut rng = Xorshift::default();
709        for _ in 0..100 {
710            rand!(rng, n: 0..200, m: 1..200, t: 0usize..=1, k: 0..[10, 1_000][t]);
711            let f = Fps998244353::from_vec(rng.random_iter(..).take(n).collect());
712            let g = Fps998244353::from_vec(rng.random_iter(..).take(m).collect());
713            let expected = f.clone().bostan_mori(g.clone(), k);
714            let result = (f * g.inv(k + 1)).data.get(k).cloned().unwrap_or_default();
715            assert_eq!(result, expected);
716
717            let f = Fps::<Modulo1000000009>::from_vec(rng.random_iter(..).take(n).collect());
718            let g = Fps::<Modulo1000000009>::from_vec(rng.random_iter(..).take(m).collect());
719            let expected = f.clone().bostan_mori(g.clone(), k);
720            let result = (f * g.inv(k + 1)).data.get(k).cloned().unwrap_or_default();
721            assert_eq!(result, expected);
722        }
723    }
724
725    #[test]
726    fn test_bostan_mori_msb() {
727        let mut rng = Xorshift::default();
728        for _ in 0..100 {
729            rand!(rng, n: 2..20, t: 0usize..=1, k: 0..[10, 1_000_000_000][t]);
730            let f = Fps998244353::from_vec(rng.random_iter(..).take(n - 1).collect());
731            let g = Fps998244353::from_vec(rng.random_iter(..).take(n).collect());
732            let expected = f.clone().bostan_mori(g.clone(), k);
733            let result = (f * g.bostan_mori_msb(k))[n - 2];
734            assert_eq!(result, expected);
735        }
736    }
737
738    #[test]
739    fn test_pow_mod() {
740        let mut rng = Xorshift::default();
741        for _ in 0..100 {
742            rand!(rng, n: 2..20, t: 0usize..=1, k: 0..[10, 1_000_000_000][t]);
743            let f = Fps998244353::from_vec(rng.random_iter(..).take(n).collect());
744            let mut expected = Fps998244353::one();
745            {
746                let mut p = Fps998244353::one() << 1;
747                let mut k = k;
748                while k > 0 {
749                    if k & 1 == 1 {
750                        expected = (expected * &p) % &f;
751                    }
752                    p = (&p * &p) % &f;
753                    k >>= 1;
754                }
755            }
756
757            let result = f.pow_mod(k);
758            assert_eq!(result, expected);
759        }
760    }
761
762    #[test]
763    fn test_sum_of_powers() {
764        let mut rng = Xorshift::default();
765        for _ in 0..100 {
766            rand!(rng, n: 0..100, m: 0..10);
767            let a: Vec<_> = rng.random_iter(..).take(n).collect();
768            let result = Fps998244353::sum_of_powers(a.iter().cloned(), m + 1);
769            for k in 0..=m {
770                let mut expected = MInt998244353::zero();
771                for &x in &a {
772                    expected += x.pow(k);
773                }
774                assert_eq!(result[k], expected);
775            }
776        }
777    }
778}
competitive/math/formal_power_series/formal_power_series_impls.rs

competitive/math/formal_power_series/
formal_power_series_impls.rs
