Day 087-Q5 — べき乗和の高速計算（連続点 Lagrange 補間 / Faulhaber の公式）

問題

正整数 $N,K$ が与えられる。$S=\sum_{i=1}^{N} i^K \pmod{10^9+7}$ を求めよ。

入出力例

3 2

14

10 1

55

例1: $1^2+2^2+3^2=14$（直接計算範囲）。例2: $1+\dots+10=55$（補間を要する範囲、Step 4で手計算検証）。

概念図: $K+2$ 点から多項式 $S(x)$ を復元して $x=N$ を評価

ヒント

# y[i] = S(i), i=0..m (m=K+1)
# prefix[i] = Π_{j<i}(N-j), suffix[i] = Π_{j>=i}(N-j)
# term_i = y[i] * prefix[i]*suffix[i+1] * inv(i!) * inv((m-i)!) * (-1)^(m-i)
# answer = Σ term_i

模範解答

import sys

MOD = 10**9 + 7

def solve():
    data = sys.stdin.buffer.read().split()
    N = int(data[0]); K = int(data[1])
    m = K + 1
    size = m + 1

    y = [0] * size
    for i in range(1, size):
        y[i] = (y[i - 1] + pow(i, K, MOD)) % MOD

    if N < size:
        print(y[N])
        return

    prefix = [1] * (size + 1)
    suffix = [1] * (size + 1)
    for i in range(size):
        prefix[i + 1] = prefix[i] * ((N - i) % MOD) % MOD
    for i in range(size - 1, -1, -1):
        suffix[i] = suffix[i + 1] * ((N - i) % MOD) % MOD

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

    ans = 0
    for i in range(size):
        num = prefix[i] * suffix[i + 1] % MOD
        denom = inv_fact[i] * inv_fact[size - 1 - i] % MOD
        term = y[i] * num % MOD * denom % MOD
        if (size - 1 - i) % 2 == 1:
            term = (-term) % MOD
        ans = (ans + term) % MOD

    print(ans % MOD)

solve()

計算量: 点の計算 $O(K\log K)$、補間 $O(K)$。全体 $O(K\log K)$、$N$ の大きさに依存しない。

Step-by-Step 解説

Step 1: $S(x)$ が多項式である理由

Faulhaber の公式より $\sum_{i=1}^x i^K$ は $x$ について次数 $K+1$ の多項式で表せる。

Step 2: サンプル点の計算

$x=0,\dots,m$（$m=K+1$）の値を累積和で計算（各点 $O(\log K)$、全体 $O(K\log K)$）。

Step 3: 連続点 Lagrange 補間

ノードが等間隔の整数のため分母 $\prod_{j\ne i}(i-j)=i!\,(m-i)!\,(-1)^{m-i}$ に閉じた式があり、分子は前から積んだ prefix と後ろから積んだ suffix の積で $O(1)$。全体 $O(K)$ で $x=N$ を評価できる。

Step 4: 手計算による検証（$N=10,K=1$）

合計 $0-80+135=55=S(10)$ と一致。

よくあるミス

次のステップ

• 発展問題: 複数の $K$ に対するクエリ（$K$ ごとに $O(K\log K)$ 前処理）
• 発展問題: 等間隔でない点集合での高速多点補間（Subproduct Tree・$O(K\log^2K)$）

自己評価