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) -> Option<Self::Poly>;
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) -> Option<Self::Poly> {
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 Some(f.into());
36        }
37        if k % 2 == 1 {
38            return None;
39        }
40
41        let x = mod_sqrt(f[k].value() as u64, P as u64)?;
42        let m = n - k;
43
44        let half = ConstModInt::new(2).inv();
45        let mut t = 1;
46        let mut ret = vec![ConstModInt::new(x as u32)];
47
48        loop {
49            let mut f = f[k..k + t.min(m)].to_vec();
50            f.resize(t, 0.into());
51
52            ret.resize(t, 0.into());
53            let h = self.fps_inv(ret.clone().into());
54            let h = self.mul(f.into(), h);
55            let h: Vec<_> = h.into();
56
57            for (x, y) in ret.iter_mut().zip(h) {
58                *x = (*x + y) * half;
59            }
60
61            if t >= m {
62                break;
63            }
64
65            t <<= 1;
66        }
67
68        ret.resize(n, 0.into());
69        Some(self.shift_higher(ret.into(), k / 2))
70    }
71}
haar_lib/math/fps/sqrt.rs

haar_lib/math/fps/
sqrt.rs
