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) -> Result<Self::Poly, &'static str>;
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) -> Result<Self::Poly, &'static str> {
18        let f: Vec<_> = f.into();
19
20        if f[0].value() == 0 {
21            return Err("定数項が`0`の形式的べき級数の逆数を計算しようとした。");
22        }
23        let n = f.len();
24
25        let mut t = 1;
26        let mut ret = vec![f[0].inv()];
27        ret.reserve(2 * n);
28
29        loop {
30            let mut f = f[0..(2 * t).min(n)].to_vec();
31            f.resize(2 * t, 0.into());
32            self.ntt.ntt(&mut f);
33
34            let mut g = ret.clone();
35            g.resize(2 * t, 0.into());
36            self.ntt.ntt(&mut g);
37
38            for (f, g) in f.iter_mut().zip(g.iter()) {
39                *f *= *g;
40            }
41            self.ntt.intt(&mut f);
42
43            let h = f;
44
45            let mut h = h[t..2 * t].to_vec();
46            h.resize(2 * t, 0.into());
47            self.ntt.ntt(&mut h);
48
49            for (h, g) in h.iter_mut().zip(g.iter()) {
50                *h *= *g;
51            }
52            self.ntt.intt(&mut h);
53
54            let g = h;
55
56            ret.resize(2 * t, 0.into());
57
58            for (ret, x) in ret.iter_mut().skip(t).zip(g.into_iter().take(t)) {
59                *ret = -x;
60            }
61
62            t <<= 1;
63
64            if t >= n {
65                break;
66            }
67        }
68
69        ret.truncate(n);
70        Ok(ret.into())
71    }
72}