haar_lib/math/fps/
pow_sparse.rs

1//! 疎な形式的冪級数の累乗
2use crate::math::polynomial::sparse::SparsePolynomial;
3use crate::math::polynomial::Polynomial;
4use crate::math::prime_mod::PrimeMod;
5use crate::num::const_modint::*;
6
7/// 疎な形式的冪級数の累乗
8pub trait FpsPowSparse {
9    /// 戻り値の型
10    type Output;
11
12    /// $f(x) = \sum_0^{n-1} a_ix^i$について、$(f(x))^m$の先頭$n$項を求める。
13    fn fps_pow_sparse(self, m: u64, n: usize) -> Result<Self::Output, &'static str>;
14}
15
16impl<P: PrimeMod> FpsPowSparse for SparsePolynomial<P> {
17    type Output = Polynomial<P>;
18
19    /// **Time complexity** $O(nk)$
20    fn fps_pow_sparse(self, m: u64, n: usize) -> Result<Self::Output, &'static str> {
21        if m == 0 {
22            let mut f: Vec<_> = vec![ConstModInt::new(0); n];
23            f[0] = ConstModInt::new(1);
24            return Ok(f.into());
25        }
26
27        let k = (0..n).find(|&i| self.coeff_of(i).value() != 0).unwrap_or(n);
28
29        if k >= n {
30            return Ok(vec![ConstModInt::new(0); n].into());
31        }
32
33        if k.checked_mul(m as usize).is_none_or(|x| x >= n) {
34            return Ok(vec![ConstModInt::new(0); n].into());
35        }
36
37        let a = self.coeff_of(k);
38
39        let mut f = self;
40        f.shift_lower(k);
41        f.scale(a.inv());
42
43        let mut g = vec![ConstModInt::new(0); n];
44        let mut ret = vec![ConstModInt::new(0); n];
45        ret[0] = 1.into();
46
47        let mut invs = vec![ConstModInt::new(1); n + 1];
48        for i in 2..=n {
49            invs[i] = -invs[P::PRIME_NUM as usize % i] * ConstModInt::new(P::PRIME_NUM / i as u32);
50        }
51
52        for i in 0..n - 1 {
53            let mut s = ConstModInt::new(0);
54
55            for (&j, &fj) in f.iter() {
56                if j != 0 {
57                    if i >= j {
58                        s -= fj * g[i - j];
59                    }
60                    if i + 1 >= j {
61                        s += ret[i - j + 1] * j.into() * fj * m.into();
62                    }
63                }
64            }
65
66            g[i] = s;
67            ret[i + 1] = s * invs[i + 1];
68        }
69
70        let mut ret = Polynomial::from(ret);
71        ret.scale(a.pow(m));
72        ret.shift_higher(m as usize * k);
73        Ok(ret)
74    }
75}
haar_lib/math/fps/pow_sparse.rs

haar_lib/math/fps/
pow_sparse.rs
