haar_lib/math/fps/
sqrt_sparse.rs

1//! 疎な形式的冪級数の平方根
2use crate::math::mod_ops::sqrt::mod_sqrt;
3use crate::math::polynomial::sparse::SparsePolynomial;
4use crate::math::polynomial::Polynomial;
5use crate::math::prime_mod::PrimeMod;
6use crate::num::const_modint::*;
7
8/// 疎な形式的冪級数の平方根
9pub trait FpsSqrtSparse {
10    /// 戻り値の型
11    type Output;
12
13    /// $f(x) = \sum_0^{n-1} a_ix^i$について、$\sqrt{f(x)}$の先頭$n$項を求める。
14    fn fps_sqrt_sparse(self, n: usize) -> Result<Self::Output, &'static str>;
15}
16
17impl<P: PrimeMod> FpsSqrtSparse for SparsePolynomial<P> {
18    type Output = Polynomial<P>;
19
20    /// **Time complexity** $O(nk)$
21    fn fps_sqrt_sparse(self, n: usize) -> Result<Self::Output, &'static str> {
22        let Some(k) = (0..n).find(|&i| self.coeff_of(i).value() != 0) else {
23            return Ok(vec![ConstModInt::new(0); n].into());
24        };
25
26        if k % 2 == 1 {
27            return Err("最小次数が偶数ではない。");
28        }
29
30        let a = self.coeff_of(k);
31        let sr = ConstModInt::new(
32            mod_sqrt(a.value() as u64, P::PRIME_NUM as u64)
33                .ok_or("最小次数項の係数に平方根が存在しない。")? as u32,
34        );
35
36        let mut f = self;
37        f.shift_lower(k);
38        f.scale(a.inv());
39
40        let mut g = vec![ConstModInt::new(0); n];
41        let mut ret = vec![ConstModInt::new(0); n];
42        ret[0] = 1.into();
43
44        let mut invs = vec![ConstModInt::new(1); n + 1];
45        for i in 2..=n {
46            invs[i] = -invs[P::PRIME_NUM as usize % i] * ConstModInt::new(P::PRIME_NUM / i as u32);
47        }
48
49        let half = ConstModInt::new(2).inv();
50
51        for i in 0..n - 1 {
52            let mut s = ConstModInt::new(0);
53
54            for (&j, &fj) in f.iter() {
55                if j != 0 {
56                    if i >= j {
57                        s -= fj * g[i - j];
58                    }
59                    if i + 1 >= j {
60                        s += ret[i - j + 1] * j.into() * fj * half;
61                    }
62                }
63            }
64
65            g[i] = s;
66            ret[i + 1] = s * invs[i + 1];
67        }
68
69        let mut ret = Polynomial::from(ret);
70        ret.scale(sr);
71        ret.shift_higher(k / 2);
72        Ok(ret)
73    }
74}
haar_lib/math/fps/sqrt_sparse.rs

haar_lib/math/fps/
sqrt_sparse.rs
