haar_lib/math/fps/
pow_sparse.rs

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

haar_lib/math/fps/
pow_sparse.rs
