Day 076-Q3 — FPS pow（多項式べき乗・exp・log・Newton法 $O(N \log N)$）

問題

形式的冪級数 $F(x) = \sum_{k=0}^{N-1} a_k x^k$（$a_0 = 1$, 係数は $\mathbb{Z}/p\mathbb{Z}$）が与えられる。

$$G(x) = F(x)^M \bmod x^N$$

を計算し、$[x^0], [x^1], \ldots, [x^{N-1}]$ の係数を出力せよ。$p = 998244353$。

入出力例

4 3
1 1 1 1

1 3 6 10

$(1+x+x^2+x^3)^3 = 1 + 3x + 6x^2 + 10x^3 + \ldots$ の最初の 4 係数。

概念図: FPS pow の計算パイプライン

$a_0 = 1$ のとき $\ln F$ が定義でき、$F^M = e^{M \ln F}$ が成立する。全体 $O(N \log N)$。

ヒント

def fps_pow(f, m, n):
    if m == 0:
        res = [0] * n; res[0] = 1; return res
    log_f = fps_log(f, n)          # O(N log N)
    m_log_f = [m * c % MOD for c in log_f]  # 各係数に M を掛ける
    return fps_exp(m_log_f, n)      # O(N log N)

def fps_log(f, n):
    df = fps_deriv(f[:n])
    inv_f = fps_inv(f, n)
    prod = fps_mul(df, inv_f, n)
    return fps_integ(prod)[:n]

def fps_exp(h, n):
    g = [1]; sz = 1
    while sz < n:
        sz <<= 1
        g = (g + [0] * (sz - len(g)))[:sz]
        log_g = fps_log(g, sz)
        tmp = [(h[i] if i < len(h) else 0) - log_g[i] for i in range(sz)]
        tmp[0] = (tmp[0] + 1) % MOD
        g = fps_mul(g, [x % MOD for x in tmp], sz)[:sz]
    return g[:n]

模範解答

import sys
input = sys.stdin.readline

MOD = 998244353
g_root = 3

def ntt(a, invert=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_root, (MOD - 1) // length, MOD)
        if invert: w = pow(w, MOD - 2, MOD)
        for i in range(0, n, length):
            wn = 1
            for k in range(length // 2):
                u = a[i + k]; v = 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 invert:
        ni = pow(n, MOD - 2, MOD)
        for i in range(n): a[i] = a[i] * ni % MOD

def fps_mul(a, b, n=None):
    rl = len(a) + len(b) - 1; sz = 1
    while sz < rl: sz <<= 1
    fa = a + [0]*(sz-len(a)); fb = b + [0]*(sz-len(b))
    ntt(fa); ntt(fb)
    fc = [fa[i]*fb[i]%MOD for i in range(sz)]
    ntt(fc, invert=True)
    return fc[:n] if n else fc

def fps_inv(f, n):
    g = [pow(f[0], MOD-2, MOD)]; sz = 1
    while sz < n:
        sz <<= 1
        g2 = (g + [0]*(sz-len(g)))
        tmp = fps_mul(f[:sz], g2, sz)
        tmp = [(2 - tmp[i]) % MOD for i in range(sz)]
        g = fps_mul(g2, tmp, sz)[:sz]
    return g[:n]

def fps_deriv(f):
    return [(i*f[i])%MOD for i in range(1, len(f))] + [0]

def fps_integ(f):
    n = len(f); inv = [0]*(n+1); inv[1] = 1
    for i in range(2, n+1): inv[i] = -(MOD//i)*inv[MOD%i]%MOD
    return [0] + [f[i-1]*inv[i]%MOD for i in range(1, n+1)]

def fps_log(f, n):
    df = fps_deriv(f[:n])
    return fps_integ(fps_mul(df, fps_inv(f, n), n))[:n]

def fps_exp(h, n):
    g = [1]; sz = 1
    while sz < n:
        sz <<= 1
        g = (g + [0]*(sz-len(g)))[:sz]
        log_g = fps_log(g, sz)
        tmp = [(h[i] if i < len(h) else 0) - log_g[i] for i in range(sz)]
        tmp[0] = (tmp[0] + 1) % MOD
        g = fps_mul(g, [x % MOD for x in tmp], sz)[:sz]
    return g[:n]

def fps_pow(f, m, n):
    if m == 0:
        res = [0]*n; res[0] = 1; return res
    log_f = fps_log(f, n)
    m_log_f = [m * c % MOD for c in log_f]
    return fps_exp(m_log_f, n)

def main():
    N, M = map(int, input().split())
    a = list(map(int, input().split()))
    a = (a + [0]*N)[:N]
    result = fps_pow(a, M % (MOD - 1), N)
    print(*result[:N])

main()

Step-by-Step 解説

Step 1: FPS の逆元（Newton 法）

$G \cdot F \equiv 1 \pmod{x^n}$ を解くために Newton 反復 $G_{k+1} = G_k(2 - F G_k)$ を使う。サイズを $1 \to 2 \to 4 \to \ldots$ と倍増させ $O(N \log N)$ 総計。

Step 2: FPS の log

$$\ln F = \int \frac{F'}{F}$$ 微分は $O(N)$、逆元は $O(N \log N)$、積は NTT で $O(N \log N)$、積分は $O(N)$。

Step 3: FPS の exp（Newton 法）

$G = \exp(H)$ を Newton 法で解く: $G_{k+1} = G_k(1 - \ln G_k + H)$。各ステップで log 計算が必要なため定数は大きいが全体 $O(N \log N)$。

Step 4: $F^M = \exp(M \ln F)$

$a_0 = 1$ なら $\ln F$ が定義できる。各係数に $M$ を掛けて exp を計算するだけ。

Step 5: Fermat の小定理による指数圧縮

$M$ が非常に大きい場合、$a^{p-1} \equiv 1 \pmod{p}$ より $M$ を $p-1 = 998244352$ で割った余りに圧縮できる（$a_0 = 1$ 前提）。

よくあるミス

次のステップ

• 発展問題: 「$a_0 = 0$ の場合の $F^M$」→ $F = x^k G$（$G_0 \ne 0$）に変換: $F^M = x^{kM} G^M$

自己評価