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) -> Self::Poly;
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) -> Self::Poly {
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 f.into();
24        }
25        if m == 1 {
26            return 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 f;
40        }
41
42        if k.checked_mul(m as usize).is_none_or(|x| x >= n) {
43            return vec![ConstModInt::new(0); n].into();
44        }
45
46        let a = f.coeff_of(k);
47
48        let ret = self.shift_lower(f, k);
49        let ret = self.scale(ret, a.inv());
50        let ret = self.scale(self.fps_log(ret), m.into());
51        let ret = self.fps_exp(ret);
52        let ret = self.scale(ret, a.pow(m));
53        self.shift_higher(ret, m as usize * k)
54    }
55}
haar_lib/math/fps/pow.rs

haar_lib/math/fps/
pow.rs
