Day 059-Q4 — 多項式Taylor Shift（f(x+c) O(NlogN)）

問題

$N-1$ 次多項式 $f(x) = \sum_{i=0}^{N-1} a_i x^i$ と整数 $c$ が与えられる。$f(x+c)$ の係数列 $b_0, b_1, \dots, b_{N-1}$ を $\mathrm{mod}\ 998244353$ で求めよ。

入出力例

4 2
1 0 1 1

11 18 7 1

$f(x) = x^3 + x^2 + 1$, $f(x+2) = (x+2)^3 + (x+2)^2 + 1$
$= x^3 + 6x^2 + 12x + 8 + x^2 + 4x + 4 + 1 = x^3 + 7x^2 + 16x + 13$
※ 実際の値で確認すること。

概念図: Taylor Shift の畳み込み変換

ヒント（段階的開示）

# 前計算
fact = [1]*(N+1)
for i in range(1,N+1): fact[i]=fact[i-1]*i%MOD
inv_fact = [1]*(N+1)
inv_fact[N] = pow(fact[N],MOD-2,MOD)
for i in range(N-1,-1,-1): inv_fact[i]=inv_fact[i+1]*(i+1)%MOD

A_hat = [a[N-1-i]*fact[N-1-i]%MOD for i in range(N)]
cp=1; B=[]
for j in range(N): B.append(cp*inv_fact[j]%MOD); cp=cp*c%MOD

conv = polymul(A_hat, B)  # NTT畳み込み
b = [conv[N-1-k]*inv_fact[k]%MOD for k in range(N)]

模範解答 (Python)

コピーimport sys
input = sys.stdin.readline

MOD = 998244353
g = 3

def ntt(a, inv=False):
    n = len(a)
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit: j ^= bit; bit >>= 1
        j ^= bit
        if i < j: a[i], a[j] = a[j], a[i]
    length = 2
    while length <= n:
        w = pow(g, (MOD-1)//length, MOD)
        if inv: w = pow(w, MOD-2, MOD)
        for i in range(0, n, length):
            wn = 1
            for k in range(length//2):
                u, v = a[i+k], a[i+k+length//2]*wn%MOD
                a[i+k] = (u+v)%MOD
                a[i+k+length//2] = (u-v)%MOD
                wn = wn*w%MOD
        length <<= 1
    if inv:
        inv_n = pow(n, MOD-2, MOD)
        for i in range(n): a[i] = a[i]*inv_n%MOD

def polymul(a, b):
    n = 1
    while n < len(a)+len(b)-1: n <<= 1
    fa = a + [0]*(n-len(a))
    fb = b + [0]*(n-len(b))
    ntt(fa); ntt(fb)
    for i in range(n): fa[i] = fa[i]*fb[i]%MOD
    ntt(fa, inv=True)
    return fa

def main():
    N, c = map(int, input().split())
    a = list(map(int, input().split()))

    fact = [1]*(N+1)
    for i in range(1, N+1): fact[i] = fact[i-1]*i%MOD
    inv_fact = [1]*(N+1)
    inv_fact[N] = pow(fact[N], MOD-2, MOD)
    for i in range(N-1, -1, -1): inv_fact[i] = inv_fact[i+1]*(i+1)%MOD

    A_hat = [a[N-1-i]*fact[N-1-i]%MOD for i in range(N)]

    cp = 1
    B = []
    for j in range(N):
        B.append(cp*inv_fact[j]%MOD)
        cp = cp*c%MOD

    conv = polymul(A_hat, B)

    b = [conv[N-1-k]*inv_fact[k]%MOD for k in range(N)]
    print(*b)

main()

Step-by-Step 解説

Step 1: Taylor Shift の公式

$$b_k = [x^k]f(x+c) = \sum_{i=k}^{N-1} a_i \binom{i}{k} c^{i-k}$$

Step 2: 畳み込みへの変換

$b_k = \frac{1}{k!}\sum_{i=k}^{N-1}(a_i \cdot i!) \cdot \frac{c^{i-k}}{(i-k)!}$。 $i' = N-1-i$ と置換すると、これは $\hat{A}[i'] = a_{N-1-i'} \cdot (N-1-i')!$ と $B[j] = c^j/j!$ の畳み込みの $[N-1-k]$ 番目の係数。

Step 3: NTT で計算

$\hat{A}$ と $B$ の畳み込みを NTT で $O(N \log N)$。

Step 4: 係数の読み出し

$b_k = \text{conv}[N-1-k] \cdot (k!)^{-1} \bmod p$。

よくあるミス

次のステップ

発展問題: $f(x+1), f(x+2), \dots, f(x+K)$ を一括で求める多点 Taylor Shift を実装せよ。

自己評価