haar_lib/math/fps/
inv.rs

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

haar_lib/math/fps/
inv.rs
