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 + PartialEq,
74{
75    pub fn trim_tail_zeros(&mut self) {
76        let mut len = self.length();
77        while len > 0 {
78            if self.data[len - 1].is_zero() {
79                len -= 1;
80            } else {
81                break;
82            }
83        }
84        self.truncate(len);
85    }
86}
87
88impl<T, C> Zero for FormalPowerSeries<T, C>
89where
90    T: PartialEq,
91{
92    fn zero() -> Self {
93        Self::from_vec(Vec::new())
94    }
95}
96impl<T, C> One for FormalPowerSeries<T, C>
97where
98    T: PartialEq + One,
99{
100    fn one() -> Self {
101        Self::from(T::one())
102    }
103}
104
105impl<T, C> IntoIterator for FormalPowerSeries<T, C> {
106    type Item = T;
107    type IntoIter = std::vec::IntoIter<T>;
108    fn into_iter(self) -> Self::IntoIter {
109        self.data.into_iter()
110    }
111}
112impl<'a, T, C> IntoIterator for &'a FormalPowerSeries<T, C> {
113    type Item = &'a T;
114    type IntoIter = Iter<'a, T>;
115    fn into_iter(self) -> Self::IntoIter {
116        self.data.iter()
117    }
118}
119impl<'a, T, C> IntoIterator for &'a mut FormalPowerSeries<T, C> {
120    type Item = &'a mut T;
121    type IntoIter = IterMut<'a, T>;
122    fn into_iter(self) -> Self::IntoIter {
123        self.data.iter_mut()
124    }
125}
126
127impl<T, C> FromIterator<T> for FormalPowerSeries<T, C> {
128    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
129        Self::from_vec(iter.into_iter().collect())
130    }
131}
132
133impl<T, C> Index<usize> for FormalPowerSeries<T, C> {
134    type Output = T;
135    fn index(&self, index: usize) -> &Self::Output {
136        &self.data[index]
137    }
138}
139impl<T, C> IndexMut<usize> for FormalPowerSeries<T, C> {
140    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
141        &mut self.data[index]
142    }
143}
144
145impl<T, C> From<T> for FormalPowerSeries<T, C> {
146    fn from(x: T) -> Self {
147        once(x).collect()
148    }
149}
150impl<T, C> From<Vec<T>> for FormalPowerSeries<T, C> {
151    fn from(data: Vec<T>) -> Self {
152        Self::from_vec(data)
153    }
154}
155
156impl<T, C> FormalPowerSeries<T, C>
157where
158    T: FormalPowerSeriesCoefficient,
159{
160    pub fn prefix_ref(&self, deg: usize) -> Self {
161        if deg < self.length() {
162            Self::from_vec(self.data[..deg].to_vec())
163        } else {
164            self.clone()
165        }
166    }
167    pub fn prefix(mut self, deg: usize) -> Self {
168        self.data.truncate(deg);
169        self
170    }
171    pub fn even(mut self) -> Self {
172        let mut keep = false;
173        self.data.retain(|_| {
174            keep = !keep;
175            keep
176        });
177        self
178    }
179    pub fn odd(mut self) -> Self {
180        let mut keep = true;
181        self.data.retain(|_| {
182            keep = !keep;
183            keep
184        });
185        self
186    }
187    pub fn diff(mut self) -> Self {
188        let mut c = T::one();
189        for x in self.iter_mut().skip(1) {
190            *x *= &c;
191            c += T::one();
192        }
193        if self.length() > 0 {
194            self.data.remove(0);
195        }
196        self
197    }
198    pub fn integral(mut self) -> Self {
199        let n = self.length();
200        self.data.insert(0, Zero::zero());
201        let mut fact = Vec::with_capacity(n + 1);
202        let mut c = T::one();
203        fact.push(c.clone());
204        for _ in 1..n {
205            fact.push(fact.last().cloned().unwrap() * c.clone());
206            c += T::one();
207        }
208        let mut invf = T::one() / (fact.last().cloned().unwrap() * c.clone());
209        for x in self.iter_mut().skip(1).rev() {
210            *x *= invf.clone() * fact.pop().unwrap();
211            invf *= c.clone();
212            c -= T::one();
213        }
214        self
215    }
216    pub fn eval(&self, x: T) -> T {
217        let mut base = T::one();
218        let mut res = T::zero();
219        for a in self.iter() {
220            res += base.clone() * a.clone();
221            base *= x.clone();
222        }
223        res
224    }
225}
226
227impl<T, C> FormalPowerSeries<T, C>
228where
229    T: FormalPowerSeriesCoefficient,
230    C: ConvolveSteps<T = Vec<T>>,
231{
232    pub fn inv(&self, deg: usize) -> Self {
233        debug_assert!(!self[0].is_zero());
234        let mut f = Self::from(T::one() / self[0].clone());
235        let mut i = 1;
236        while i < deg {
237            let g = self.prefix_ref((i * 2).min(deg));
238            let h = f.clone();
239            let mut g = C::transform(g.data, 2 * i);
240            let h = C::transform(h.data, 2 * i);
241            C::multiply(&mut g, &h);
242            let mut g = Self::from_vec(C::inverse_transform(g, 2 * i));
243            g >>= i;
244            let mut g = C::transform(g.data, 2 * i);
245            C::multiply(&mut g, &h);
246            let g = Self::from_vec(C::inverse_transform(g, 2 * i));
247            f.data.extend((-g).into_iter().take(i));
248            i *= 2;
249        }
250        f.truncate(deg);
251        f
252    }
253    pub fn exp(&self, deg: usize) -> Self {
254        debug_assert!(self[0].is_zero());
255        let mut f = Self::one();
256        let mut i = 1;
257        while i < deg {
258            let mut g = -f.log(i * 2);
259            g[0] += T::one();
260            for (g, x) in g.iter_mut().zip(self.iter().take(i * 2)) {
261                *g += x.clone();
262            }
263            f = (f * g).prefix(i * 2);
264            i *= 2;
265        }
266        f.prefix(deg)
267    }
268    pub fn log(&self, deg: usize) -> Self {
269        (self.inv(deg) * self.clone().diff()).integral().prefix(deg)
270    }
271    pub fn pow(&self, rhs: usize, deg: usize) -> Self {
272        if rhs == 0 {
273            return Self::from_vec(
274                once(T::one())
275                    .chain(repeat_with(T::zero))
276                    .take(deg)
277                    .collect(),
278            );
279        }
280        if let Some(k) = self.iter().position(|x| !x.is_zero()) {
281            if k >= (deg + rhs - 1) / rhs {
282                Self::zeros(deg)
283            } else {
284                let mut x0 = self[k].clone();
285                let rev = T::one() / x0.clone();
286                let x = {
287                    let mut x = T::one();
288                    let mut y = rhs;
289                    while y > 0 {
290                        if y & 1 == 1 {
291                            x *= x0.clone();
292                        }
293                        x0 *= x0.clone();
294                        y >>= 1;
295                    }
296                    x
297                };
298                let mut f = (self.clone() * &rev) >> k;
299                f = (f.log(deg) * &T::from(rhs)).exp(deg) * &x;
300                f.truncate(deg - k * rhs);
301                f <<= k * rhs;
302                f
303            }
304        } else {
305            Self::zeros(deg)
306        }
307    }
308}
309
310impl<T, C> FormalPowerSeries<T, C>
311where
312    T: FormalPowerSeriesCoefficientSqrt,
313    C: ConvolveSteps<T = Vec<T>>,
314{
315    pub fn sqrt(&self, deg: usize) -> Option<Self> {
316        if self[0].is_zero() {
317            if let Some(k) = self.iter().position(|x| !x.is_zero()) {
318                if k % 2 != 0 {
319                    return None;
320                } else if deg > k / 2 {
321                    return Some((self >> k).sqrt(deg - k / 2)? << (k / 2));
322                }
323            }
324        } else {
325            let inv2 = T::one() / (T::one() + T::one());
326            let mut f = Self::from(self[0].sqrt_coefficient()?);
327            let mut i = 1;
328            while i < deg {
329                f = (&f + &(self.prefix_ref(i * 2) * f.inv(i * 2))).prefix(i * 2) * &inv2;
330                i *= 2;
331            }
332            f.truncate(deg);
333            return Some(f);
334        }
335        Some(Self::zeros(deg))
336    }
337}
338
339impl<T, C> FormalPowerSeries<T, C>
340where
341    T: FormalPowerSeriesCoefficient,
342    C: ConvolveSteps<T = Vec<T>>,
343{
344    pub fn count_subset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
345    where
346        F: FnMut(usize) -> T,
347    {
348        let n = self.length();
349        let mut f = Self::zeros(n);
350        for i in 1..n {
351            if !self[i].is_zero() {
352                for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
353                    if j & 1 != 0 {
354                        f[d] += self[i].clone() * &inverse(j);
355                    } else {
356                        f[d] -= self[i].clone() * &inverse(j);
357                    }
358                }
359            }
360        }
361        f.exp(deg)
362    }
363    pub fn count_multiset_sum<F>(&self, deg: usize, mut inverse: F) -> Self
364    where
365        F: FnMut(usize) -> T,
366    {
367        let n = self.length();
368        let mut f = Self::zeros(n);
369        for i in 1..n {
370            if !self[i].is_zero() {
371                for (j, d) in (0..n).step_by(i).enumerate().skip(1) {
372                    f[d] += self[i].clone() * &inverse(j);
373                }
374            }
375        }
376        f.exp(deg)
377    }
378    pub fn bostan_mori(self, rhs: Self, mut n: usize) -> T {
379        let mut p = self;
380        let mut q = rhs;
381        while n > 0 {
382            let mut mq = q.clone();
383            mq.iter_mut()
384                .skip(1)
385                .step_by(2)
386                .for_each(|x| *x = -x.clone());
387            let u = p * mq.clone();
388            p = if n % 2 == 0 { u.even() } else { u.odd() };
389            q = (q * mq).even();
390            n /= 2;
391        }
392        p[0].clone() / q[0].clone()
393    }
394    fn middle_product(self, other: &C::F, deg: usize) -> Self {
395        let n = self.length();
396        let mut s = C::transform(self.reversed().data, deg);
397        C::multiply(&mut s, other);
398        Self::from_vec((C::inverse_transform(s, deg))[n - 1..].to_vec())
399    }
400    pub fn multipoint_evaluation(self, points: &[T]) -> Vec<T> {
401        let n = points.len();
402        if n <= 32 {
403            return points.iter().map(|p| self.eval(p.clone())).collect();
404        }
405        let mut subproduct_tree = Vec::with_capacity(n * 2);
406        subproduct_tree.resize_with(n, Zero::zero);
407        for x in points {
408            subproduct_tree.push(Self::from_vec(vec![-x.clone(), T::one()]));
409        }
410        for i in (1..n).rev() {
411            subproduct_tree[i] = &subproduct_tree[i * 2] * &subproduct_tree[i * 2 + 1];
412        }
413        let mut uptree_t = Vec::with_capacity(n * 2);
414        uptree_t.resize_with(1, Zero::zero);
415        subproduct_tree.reverse();
416        subproduct_tree.pop();
417        let m = self.length();
418        let v = subproduct_tree.pop().unwrap().reversed().resized(m);
419        let s = C::transform(self.data, m * 2);
420        uptree_t.push(v.inv(m).middle_product(&s, m * 2).resized(n).reversed());
421        for i in 1..n {
422            let subl = subproduct_tree.pop().unwrap();
423            let subr = subproduct_tree.pop().unwrap();
424            let (dl, dr) = (subl.length(), subr.length());
425            let len = dl.max(dr) + uptree_t[i].length();
426            let s = C::transform(uptree_t[i].data.to_vec(), len);
427            uptree_t.push(subr.middle_product(&s, len).prefix(dl));
428            uptree_t.push(subl.middle_product(&s, len).prefix(dr));
429        }
430        uptree_t[n..]
431            .iter()
432            .map(|u| u.data.first().cloned().unwrap_or_else(Zero::zero))
433            .collect()
434    }
435    pub fn product_all<I>(iter: I, deg: usize) -> Self
436    where
437        I: IntoIterator<Item = Self>,
438    {
439        let mut heap: BinaryHeap<_> = iter
440            .into_iter()
441            .map(|f| PartialIgnoredOrd(Reverse(f.length()), f))
442            .collect();
443        while let Some(PartialIgnoredOrd(_, x)) = heap.pop() {
444            if let Some(PartialIgnoredOrd(_, y)) = heap.pop() {
445                let z = (x * y).prefix(deg);
446                heap.push(PartialIgnoredOrd(Reverse(z.length()), z));
447            } else {
448                return x;
449            }
450        }
451        Self::one()
452    }
453    pub fn sum_all_rational<I>(iter: I, deg: usize) -> (Self, Self)
454    where
455        I: IntoIterator<Item = (Self, Self)>,
456    {
457        let mut heap: BinaryHeap<_> = iter
458            .into_iter()
459            .map(|(f, g)| PartialIgnoredOrd(Reverse(f.length().max(g.length())), (f, g)))
460            .collect();
461        while let Some(PartialIgnoredOrd(_, (xa, xb))) = heap.pop() {
462            if let Some(PartialIgnoredOrd(_, (ya, yb))) = heap.pop() {
463                let zb = (&xb * &yb).prefix(deg);
464                let za = (xa * yb + ya * xb).prefix(deg);
465                heap.push(PartialIgnoredOrd(
466                    Reverse(za.length().max(zb.length())),
467                    (za, zb),
468                ));
469            } else {
470                return (xa, xb);
471            }
472        }
473        (Self::zero(), Self::one())
474    }
475    pub fn kth_term_of_linearly_recurrence(self, a: Vec<T>, k: usize) -> T {
476        if let Some(x) = a.get(k) {
477            return x.clone();
478        }
479        let p = (Self::from_vec(a).prefix(self.length() - 1) * &self).prefix(self.length() - 1);
480        p.bostan_mori(self, k)
481    }
482    pub fn kth_term(a: Vec<T>, k: usize) -> T {
483        if let Some(x) = a.get(k) {
484            return x.clone();
485        }
486        Self::from_vec(berlekamp_massey(&a)).kth_term_of_linearly_recurrence(a, k)
487    }
488    /// sum_i a_i exp(b_i x)
489    pub fn linear_sum_of_exp<I, F>(iter: I, deg: usize, mut inv_fact: F) -> Self
490    where
491        I: IntoIterator<Item = (T, T)>,
492        F: FnMut(usize) -> T,
493    {
494        let (p, q) = Self::sum_all_rational(
495            iter.into_iter()
496                .map(|(a, b)| (Self::from_vec(vec![a]), Self::from_vec(vec![T::one(), -b]))),
497            deg,
498        );
499        let mut f = (p * q.inv(deg)).prefix(deg);
500        for i in 0..f.length() {
501            f[i] *= inv_fact(i);
502        }
503        f
504    }
505}
506
507impl<M, C> FormalPowerSeries<MInt<M>, C>
508where
509    M: MIntConvert<usize>,
510    C: ConvolveSteps<T = Vec<MInt<M>>>,
511{
512    /// f(x) <- f(x + a)
513    pub fn taylor_shift(mut self, a: MInt<M>, f: &MemorizedFactorial<M>) -> Self {
514        let n = self.length();
515        for i in 0..n {
516            self.data[i] *= f.fact[i];
517        }
518        self.data.reverse();
519        let mut b = a;
520        let mut g = Self::from_vec(f.inv_fact[..n].to_vec());
521        for i in 1..n {
522            g[i] *= b;
523            b *= a;
524        }
525        self *= g;
526        self.truncate(n);
527        self.data.reverse();
528        for i in 0..n {
529            self.data[i] *= f.inv_fact[i];
530        }
531        self
532    }
533}
competitive/math/formal_power_series/formal_power_series_impls.rs

competitive/math/formal_power_series/
formal_power_series_impls.rs
