haar_lib/math/fps/
sqrt_sparse.rs

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

haar_lib/math/fps/
sqrt_sparse.rs
