haar_lib/math/polynomial/
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::math::prime_mod::PrimeMod;
8use crate::num::const_modint::*;
9
10/// Polynomial Taylor shift
11pub trait TaylorShift {
12    /// 多項式の係数の型
13    type Value;
14
15    /// 多項式 `p` = $f(x) = a_0 + a_1x + \cdots + a_nx^n$に対して、<br>
16    /// 多項式 $f(x + c) = a_0 + a_1(x + c) + \cdots + a_n(x + c)^n = b_0 + b_0x + \cdots + b_nx^n$
17    /// を満たす、数列{$b_i$}を求める。
18    fn taylor_shift(self, c: Self::Value) -> Self;
19}
20
21impl<P: PrimeMod> TaylorShift for Polynomial<P> {
22    type Value = ConstModInt<P>;
23
24    fn taylor_shift(self, c: Self::Value) -> Self {
25        let p: Vec<_> = self.into();
26        let n = p.len();
27        let mut f = ConstModInt::new(1);
28
29        let mut a = vec![ConstModInt::new(0); 2 * n - 1];
30        for (i, (a, p)) in a.iter_mut().skip(n - 1).zip(p.into_iter()).enumerate() {
31            if i != 0 {
32                f *= ConstModInt::new(i as u32);
33            }
34            *a = p * f;
35        }
36
37        let mut g = vec![ConstModInt::new(0); n];
38        g[n - 1] = f.inv();
39        for i in (0..n - 1).rev() {
40            g[i] = g[i + 1] * ConstModInt::new(i as u32 + 1);
41        }
42
43        let mut d = ConstModInt::new(1);
44        let mut b = vec![ConstModInt::new(0); 2 * n - 1];
45        for (b, g) in b.iter_mut().take(n).rev().zip(g.iter()) {
46            *b = d * *g;
47            d *= c;
48        }
49
50        let c = Self::NTT.convolve(a, b);
51        c.into_iter()
52            .skip((n - 1) * 2)
53            .zip(g)
54            .map(|(c, g)| c * g)
55            .collect::<Vec<_>>()
56            .into()
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use crate::math::prime_mod::Prime;
63
64    use super::*;
65
66    type P = Prime<998244353>;
67
68    #[test]
69    fn test() {
70        let coeffs = vec![1, 2, 3, 4, 5];
71        let c = 3;
72
73        let res = Polynomial::<P>::from(coeffs.clone()).taylor_shift(c.into());
74
75        let mut ans = Polynomial::zero();
76        for (i, a) in coeffs.into_iter().enumerate() {
77            let mut p = Polynomial::<P>::from(vec![c, 1]).pow(i as u64);
78            p.scale(a.into());
79            ans += p;
80        }
81
82        assert_eq!(res, ans);
83    }
84}
haar_lib/math/polynomial/polynomial_taylor_shift.rs

haar_lib/math/polynomial/
polynomial_taylor_shift.rs
