haar_lib/math/
polynomial_taylor_shift.rs

1//! 多項式$f(x)$に対して、$f(x + c)$の係数を求める。
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/polynomial_taylor_shift>
5
6use crate::math::polynomial::*;
7use crate::num::const_modint::*;
8
9/// Polynomial Taylor shift
10pub trait TaylorShift {
11    /// 多項式の型
12    type Poly;
13    /// 多項式の係数の型
14    type Value;
15
16    /// 多項式 `p` = $f(x) = a_0 + a_1x + \cdots + a_nx^n$に対して、<br>
17    /// 多項式 $f(x + c) = a_0 + a_1(x + c) + \cdots + a_n(x + c)^n = b_0 + b_0x + \cdots + b_nx^n$
18    /// を満たす、数列{$b_i$}を求める。
19    fn taylor_shift(&self, p: Self::Poly, c: Self::Value) -> Self::Poly;
20}
21
22impl<const P: u32, const PR: u32> TaylorShift for PolynomialOperator<'_, P, PR> {
23    type Poly = Polynomial<P>;
24    type Value = ConstModInt<P>;
25
26    fn taylor_shift(&self, p: Self::Poly, c: Self::Value) -> Self::Poly {
27        let p: Vec<_> = p.into();
28        let n = p.len();
29        let mut f = ConstModInt::new(1);
30
31        let mut a = vec![ConstModInt::new(0); 2 * n - 1];
32        for (i, (a, p)) in a.iter_mut().skip(n - 1).zip(p.into_iter()).enumerate() {
33            if i != 0 {
34                f *= ConstModInt::new(i as u32);
35            }
36            *a = p * f;
37        }
38
39        let mut g = vec![ConstModInt::new(0); n];
40        g[n - 1] = f.inv();
41        for i in (0..n - 1).rev() {
42            g[i] = g[i + 1] * ConstModInt::new(i as u32 + 1);
43        }
44
45        let mut d = ConstModInt::new(1);
46        let mut b = vec![ConstModInt::new(0); 2 * n - 1];
47        for (b, g) in b.iter_mut().take(n).rev().zip(g.iter()) {
48            *b = d * *g;
49            d *= c;
50        }
51
52        //    let c = ntt.convolve(a, b);
53        let c: Vec<_> = self.mul(a.into(), b.into()).into();
54        c.into_iter()
55            .skip((n - 1) * 2)
56            .zip(g)
57            .map(|(c, g)| c * g)
58            .collect::<Vec<_>>()
59            .into()
60    }
61}
haar_lib/math/polynomial_taylor_shift.rs

haar_lib/math/
polynomial_taylor_shift.rs
