haar_lib/math/fps/
sqrt.rs

1//! 形式的冪級数の平方根
2use crate::num::ff::*;
3use crate::{
4    math::{
5        fps::inv::FpsInv,
6        mod_ops::sqrt::mod_sqrt,
7        polynomial::{Polynomial, PolynomialOperator},
8    },
9    num::const_modint::ConstModInt,
10};
11
12/// 形式的冪級数の平方根
13pub trait FpsSqrt {
14    /// 多項式の型
15    type Poly;
16
17    /// $f(x) = \sum_0^{n-1} a_ix^i$について、$\sqrt{f(x)}$の先頭$n$項を求める。
18    fn fps_sqrt(&self, f: Self::Poly) -> Result<Self::Poly, &'static str>;
19}
20
21impl<const P: u32, const PR: u32> FpsSqrt for PolynomialOperator<'_, P, PR> {
22    type Poly = Polynomial<P>;
23
24    fn fps_sqrt(&self, f: Self::Poly) -> Result<Self::Poly, &'static str> {
25        let f: Vec<_> = f.into();
26
27        let n = f.len();
28        let k = f
29            .iter()
30            .enumerate()
31            .find(|(_, &x)| x.value() != 0)
32            .map_or(n, |(k, _)| k);
33
34        if k == n {
35            return Ok(f.into());
36        }
37        if k % 2 == 1 {
38            return Err("最小次数が偶数ではない。");
39        }
40
41        let x = mod_sqrt(f[k].value() as u64, P as u64)
42            .ok_or("最小次数項の係数に平方根が存在しない。")?;
43        let m = n - k;
44
45        let half = ConstModInt::new(2).inv();
46        let mut t = 1;
47        let mut ret = vec![ConstModInt::new(x as u32)];
48
49        loop {
50            let mut f = f[k..k + t.min(m)].to_vec();
51            f.resize(t, 0.into());
52
53            ret.resize(t, 0.into());
54            let h = self.fps_inv(ret.clone().into())?;
55            let h = self.mul(f.into(), h);
56            let h: Vec<_> = h.into();
57
58            for (x, y) in ret.iter_mut().zip(h) {
59                *x = (*x + y) * half;
60            }
61
62            if t >= m {
63                break;
64            }
65
66            t <<= 1;
67        }
68
69        ret.resize(n, 0.into());
70        let mut ret: Polynomial<P> = ret.into();
71        ret.shift_higher(k / 2);
72        Ok(ret)
73    }
74}
haar_lib/math/fps/sqrt.rs

haar_lib/math/fps/
sqrt.rs
