haar_lib/math/fps/
sqrt.rs

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

haar_lib/math/fps/
sqrt.rs
