Lagrange 補間

$x_0, \ldots, x_d$ が互いに異なるとき、点 $(x_i, y_i)$ をすべて通る $d$ 次以下の多項式は $f(x) = \sum_i y_i \prod_{j \ne i} \frac{x - x_j}{x_i - x_j}$ で、これを mod で計算します。 $x_i = i$ のときは、分子を左右からの累積積、分母を階乗で表せるので $O(d)$ になります。式の導き方は 多項式補間についての記事 に書いています。

mod は素数で、 $x$ は大きくても負でもかまいません。

• lagrange(xs, ys, x, mod): 点 $(x_i, y_i)$ を通る多項式の $f(x)$。 $x_i$ は mod で互いに異なる。 $O(d^2 + d \log \mathrm{mod})$
• lagrange_iota(ys, x, mod): 点 $(i, y_i)$（ $i = 0, 1, \ldots, d$）を通る多項式の $f(x)$。 $d < \mathrm{mod}$。 $O(d + \log \mathrm{mod})$

答えが $n$ についての $d$ 次多項式になるとわかっているときに、 $n = 0, 1, \ldots, d$ の答えを愚直に求めておき、大きい $n$ での値を求めるのによく使います。

例（ $1^k + 2^k + \cdots + n^k$ は $n$ の $k + 1$ 次多項式）:

mod = 998244353
ys = [0]
for i in range(1, k + 2):
    ys.append((ys[-1] + pow(i, k, mod)) % mod)
print(lagrange_iota(ys, n, mod))

コード

def lagrange(xs, ys, x, mod):
    # 点 (xs[i], ys[i]) をすべて通る len(xs) - 1 次以下の多項式 f について f(x) を mod で求める
    # xs は mod で互いに異なり、mod は素数
    n = len(xs)
    res = 0
    for i in range(n):
        num = 1
        den = 1
        for j in range(n):
            if i == j:
                continue
            num = num * (x - xs[j]) % mod
            den = den * (xs[i] - xs[j]) % mod
        res += ys[i] * num * pow(den, mod - 2, mod)
    return res % mod


def lagrange_iota(ys, x, mod):
    # xs = 0, 1, ..., d のとき（ys[i] = f(i)）。mod は d より大きい素数
    d = len(ys) - 1
    x %= mod
    if x <= d:
        return ys[x] % mod
    pre = [1] * (d + 2)  # pre[i] = (x - 0)(x - 1)...(x - (i - 1))
    for i in range(d + 1):
        pre[i + 1] = pre[i] * (x - i) % mod
    suf = [1] * (d + 2)  # suf[i] = (x - i)(x - (i + 1))...(x - d)
    for i in range(d, -1, -1):
        suf[i] = suf[i + 1] * (x - i) % mod
    fact = [1] * (d + 1)
    for i in range(1, d + 1):
        fact[i] = fact[i - 1] * i % mod
    inv_fact = [1] * (d + 1)
    inv_fact[d] = pow(fact[d], mod - 2, mod)
    for i in range(d, 0, -1):
        inv_fact[i - 1] = inv_fact[i] * i % mod
    res = 0
    for i in range(d + 1):
        t = ys[i] * pre[i] % mod * suf[i + 1] % mod * inv_fact[i] % mod * inv_fact[d - i] % mod
        res += -t if (d - i) & 1 else t
    return res % mod