Polynomial Taylor Shift を履修したので、簡単にまとめてみます。例によって、Yosupo Library Checker の問題を通します。

なお、Polynomial Taylor Shift の解法は、単なるライブラリ整備の一環とみなすというよりは、その導出過程をぜひ押さえておきたい問題だと思いました。他の問題でも使える典型的な考え方がたくさん登場します！

1. Polynomial Taylor Shift とは

Polynomial Taylor Shift とは、整数係数の多項式

$f(x) = \displaystyle \sum_{i=0}^{N} a_{i} x^{i}$

が与えられたときに、整数定数 $c$ に対して $f(x+c)$ を展開したときの各係数を求めたいという問題です。ここでは、mod 998244353 で求めます。

たとえば、

$f(x) = 1 + 2x + 3x^{2} + 4x^{3} + 5x^{4}$
$c = 3$

のとき、

$f(x+c) = 547 + 668x + 309x^{2} + 64x^{3} + 5x^{4}$

となります。

愚直に計算すると $O(N^{2})$ 個の項を整理することになりますので、 $O(N^{2})$ の計算量が必要になります。しかし実は、畳み込み計算一発で求めることができて、 $O(N \log N)$ の計算量で求められます。

2. 前提知見

Polynomial Taylor Shift の解法に登場する知見を 2 つ紹介します。これらの知見は、Polynomial Taylor Shift 以外のさまざまな問題でも使える典型的・汎用的なものです！

知見その 1：添字の差が一定の積和 → 畳み込み

畳み込みというと、添字の和が一定の積和を計算するイメージが強いです。具体的には、2 つの数列 $a_{0}, a_{1}, \dots, a_{N}$ 、 $b_{0}, b_{1}, \dots, b_{N}$ に対して、

$c_{i} = \displaystyle \sum_{j=0}^{i} a_{j} b_{i-j}$

によって定まる数列 $\{c_{i}\}$ を $O(N \log N)$ の計算量で高速に計算することができます。多項式で表すならば、

$p(x) = a_{0} + a_{1} x + a_{2} x^{2} + \dots + a_{N-1} x^{N-1} + a_{N} x^{N}$
$q(x) = b_{0} + b_{1} x + b_{2} x^{2} + \dots + b_{N-1} x^{N-1} + b_{N} x^{N}$

とすると次のように表せます。

$c_{i} = \lbrack x^{i} \rbrack p(x)q(x)$

一方、実は、添字の差が一定の積和を計算する際にも、高速畳み込みアルゴリズムを活用できます。具体的には、

$c_{i} = \displaystyle \sum_{j=0}^{N-i} a_{j+i} b_{j}$

を畳み込み計算に帰着してみましょう。2 つの多項式 $p(x), q(x)$ を次のように定義します。 $q(x)$ の係数が先程とは逆になっています。

$p(x) = a_{0} + a_{1} x + a_{2} x^{2} + \dots + a_{N-1} x^{N-1} + a_{N} x^{N}$
$q(x) = b_{N} + b_{N-1} x + b_{N-2} x^{2} + \dots + b_{1} x^{N-1} + b_{0} x^{N}$

このとき、 $a_{j+i}$ を係数にもつ項の次数は $j+i$ 、 $b_{j}$ を係数にもつ項の次数は $N-j$ なので、

$c_{i} = \lbrack x^{i+N} \rbrack p(x)q(x)$

となります。つまり、多項式の積 $p(x)q(x)$ を NTT などで計算することによって、 $c_{0}, c_{1}, \dots, c_{N}$ の値を計算量 $O(N \log N)$ で求められます。

知見その 2：シグマの入れ替え

式変形をしていると、次の形の二重シグマにしばしば出会います。

$\displaystyle \sum_{i=0}^{N} \sum_{j=0}^{i} f(i, j)$

これは次のようにして、添字の順序を入れ替えることができます。

$\displaystyle \sum_{i=0}^{N} \sum_{j=0}^{i} f(i, j) = \sum_{j=0}^{N} \sum_{i=j}^{N} f(i, j)$

これは一体どうなっているのでしょうか。一般に、二重シグマ計算で添字を入れ替えたいとき、添字の組 $(i, j)$ の表す領域を二次元平面上に図示することが有効です。 $N = 5$ のとき、下図のようになります。

上の式変形は、左側の式は下図の左側のように縦方向に見ていて、右側の式は下図の右側のように横方向に見ていると解釈できます。

このような式変形は、競プロをしていると、さまざまな場面で登場します！

3. Polynomial Taylor Shift の解法

これらの前提知見を踏まえて、Polynomial Taylor Shift を畳み込みへと帰着します。 $f(x) = \displaystyle \sum_{i=0}^{N} a_{i} x^{i}$ に対して、 $f(x+c)$ を二項定理を用いて計算していきましょう。

$f(x+c)$
$= \displaystyle \sum_{i=0}^{N} a_{i} (x+c)^{i}$
$= \displaystyle \sum_{i=0}^{N} a_{i} \sum_{j=0}^{i} \binom{i}{j} c^{i - j} x^{j}$
$= \displaystyle \sum_{i=0}^{N} \sum_{j=0}^{i} (a_{i} i!) \frac{c^{i - j}}{(i - j)!} \frac{x^{j}}{j!}$

ここで、「知見その 2：シグマの入れ替え」を用いて、

$f(x+c)$
$= \displaystyle \sum_{j=0}^{N} \sum_{i=j}^{N} (a_{i} i!) \frac{c^{i - j}}{(i - j)!} \frac{x^{j}}{j!}$
$= \displaystyle \sum_{j=0}^{N} \frac{x^{j}}{j!} \sum_{i=j}^{N} (a_{i} i!) \frac{c^{i - j}}{(i - j)!}$
$= \displaystyle \sum_{j=0}^{N} \frac{x^{j}}{j!} \sum_{i=0}^{N-j} (a_{i+j} (i+j)!) \frac{c^{i}}{i!}$

となります。ここで、最後の式変形は、添字 $i$ が 0 始まりとなるように調整しました。これより、 $f(x+c)$ の各係数 $b_{0}, b_{1}, \dots, b_{N}$ は、次のように表せることがわかりました。

$b_{j} = \displaystyle \frac{1}{j!} \sum_{i=0}^{N-j} (a_{i+j} (i+j)!) \frac{c^{i}}{i!}$

さて、このシグマ計算の部分は、まさに「知見その 1：添字の差が一定の積和 → 畳み込み」を適用できる形となっています。よって、全体として $O(N \log N)$ の計算量で $f(x + c)$ の各係数値が求められることが分かりました。

4. 実装例

ここでは、Yosupo Library Checker の Polynomial Taylor Shift を通るコードを書きます。ACL の convolution を活用しますた。

#include <bits/stdc++.h>
#include "atcoder/convolution.hpp"
#include "atcoder/modint.hpp"
using namespace std;
using namespace atcoder;
using mint = modint998244353;

// Binomial coefficient
const int MOD = 998244353;
template<class T> struct BiCoef {
    vector<T> fact_, inv_, finv_;
    constexpr BiCoef() {}
    constexpr BiCoef(int n) noexcept : fact_(n, 1), inv_(n, 1), finv_(n, 1) {
        init(n);
    }
    constexpr void init(int n) noexcept {
        fact_.assign(n, 1), inv_.assign(n, 1), finv_.assign(n, 1);
        for(int i = 2; i < n; i++){
            fact_[i] = fact_[i-1] * i;
            inv_[i] = -inv_[MOD%i] * (MOD/i);
            finv_[i] = finv_[i-1] * inv_[i];
        }
    }
    constexpr T com(int n, int k) const noexcept {
        if (n < k || n < 0 || k < 0) return 0;
        return fact_[n] * finv_[k] * finv_[n-k];
    }
    constexpr T fact(int n) const noexcept {
        if (n < 0) return 0;
        return fact_[n];
    }
    constexpr T inv(int n) const noexcept {
        if (n < 0) return 0;
        return inv_[n];
    }
    constexpr T finv(int n) const noexcept {
        if (n < 0) return 0;
        return finv_[n];
    }
};

// Polynomial Taylor Shift
// given: f(x), c
// find: coefficientss of f(x + c)
template<class mint> vector<mint> PolynomialTaylorShift(const vector<mint> &f, long long c) {
    int N = (int)f.size() - 1;
    BiCoef<mint> bc(N + 1);
    
    // convollution
    vector<mint> p(N + 1), q(N + 1);
    for (int i = 0; i <= N; ++i) {
        p[i] = f[i] * bc.fact(i);
        q[N - i] = mint(c).pow(i) * bc.finv(i);
    }
    vector<mint> pq = convolution(p, q);
    
    // find
    vector<mint> res(N + 1);
    for (int i = 0; i <= N; ++i) res[i] = pq[i + N] * bc.finv(i);
    return res;
}

int main() {
    long long N, c;
    cin >> N >> c;
    vector<mint> a(N);
    for (int i = 0; i < N; ++i) {
        long long v;
        cin >> v;
        a[i] = v;
    }
    
    vector<mint> res = PolynomialTaylorShift(a, c);
    for (int i = 0; i < N; ++i) cout << res[i].val() << " ";
    cout << endl;
}