haar_lib/math/fps/
pow.rs

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

haar_lib/math/fps/
pow.rs
