haar_lib/math/fps/
pow.rs

1//! 形式的冪級数の累乗
2use crate::math::fps::{exp::*, log::*};
3use crate::math::polynomial::{Polynomial, PolynomialOperator};
4use crate::num::{const_modint::ConstModInt, ff::*};
5
6/// 形式的冪級数の累乗
7pub trait FpsPow {
8    /// 多項式の型
9    type Poly;
10
11    /// $f(x) = \sum_0^{n-1} a_ix^i$について、$(f(x))^m$の先頭$n$項を求める。
12    fn fps_pow(&self, f: Self::Poly, m: u64) -> Result<Self::Poly, &'static str>;
13}
14
15impl<const P: u32, const PR: u32> FpsPow for PolynomialOperator<'_, P, PR> {
16    type Poly = Polynomial<P>;
17
18    fn fps_pow(&self, f: Self::Poly, m: u64) -> Result<Self::Poly, &'static str> {
19        if m == 0 {
20            let mut f: Vec<_> = f.into();
21            f.fill(ConstModInt::new(0));
22            f[0] = ConstModInt::new(1);
23            return Ok(f.into());
24        }
25        if m == 1 {
26            return Ok(f);
27        }
28
29        let n = f.len();
30        let mut k = 0;
31        while k < n {
32            if f.coeff_of(k).value() != 0 {
33                break;
34            }
35            k += 1;
36        }
37
38        if k >= n {
39            return Ok(f);
40        }
41
42        if k.checked_mul(m as usize).is_none_or(|x| x >= n) {
43            return Ok(vec![ConstModInt::new(0); n].into());
44        }
45
46        let a = f.coeff_of(k);
47
48        let mut ret = f;
49        ret.shift_lower(k);
50        ret.scale(a.inv());
51        let mut ret = self.fps_log(ret)?;
52        ret.scale(m.into());
53        let mut ret = self.fps_exp(ret)?;
54        ret.scale(a.pow(m));
55        ret.shift_higher(m as usize * k);
56        Ok(ret)
57    }
58}
haar_lib/math/fps/pow.rs

haar_lib/math/fps/
pow.rs
