haar_lib/math/fps/
inv.rs

1//! 形式的冪級数の逆数
2use crate::math::polynomial::{Polynomial, PolynomialOperator};
3use crate::num::ff::*;
4
5/// 形式的冪級数の逆数
6pub trait FpsInv {
7    /// 多項式の型
8    type Poly;
9
10    /// $f(x) = \sum_0^{n-1} a_ix^i$について、$\frac{1}{f(x)}$の先頭$n$項を求める。
11    fn fps_inv(&self, f: Self::Poly) -> Self::Poly;
12}
13
14impl<const P: u32, const PR: u32> FpsInv for PolynomialOperator<'_, P, PR> {
15    type Poly = Polynomial<P>;
16
17    fn fps_inv(&self, f: Self::Poly) -> Self::Poly {
18        let f: Vec<_> = f.into();
19        assert_ne!(f[0].value(), 0);
20        let n = f.len();
21
22        let mut t = 1;
23        let mut ret = vec![f[0].inv()];
24        ret.reserve(2 * n);
25
26        loop {
27            let mut f = f[0..(2 * t).min(n)].to_vec();
28            f.resize(2 * t, 0.into());
29            self.ntt.ntt(&mut f);
30
31            let mut g = ret.clone();
32            g.resize(2 * t, 0.into());
33            self.ntt.ntt(&mut g);
34
35            for (f, g) in f.iter_mut().zip(g.iter()) {
36                *f *= *g;
37            }
38            self.ntt.intt(&mut f);
39
40            let h = f;
41
42            let mut h = h[t..2 * t].to_vec();
43            h.resize(2 * t, 0.into());
44            self.ntt.ntt(&mut h);
45
46            for (h, g) in h.iter_mut().zip(g.iter()) {
47                *h *= *g;
48            }
49            self.ntt.intt(&mut h);
50
51            let g = h;
52
53            ret.resize(2 * t, 0.into());
54
55            for (ret, x) in ret.iter_mut().skip(t).zip(g.into_iter().take(t)) {
56                *ret = -x;
57            }
58
59            t <<= 1;
60
61            if t >= n {
62                break;
63            }
64        }
65
66        ret.truncate(n);
67        ret.into()
68    }
69}
haar_lib/math/fps/inv.rs

haar_lib/math/fps/
inv.rs
