ABC 137 F - Polynomial Construction

解法

想定解はかなり頭がいいが、今回は汎用的なラグランジュ補間で解く。
今 \(n\) 次の多項式 \(P(x)\) が \(P(x_i) = y_i ~(i = 0, \ldots, n)\) を満たすとする。このとき、$$P(x) = \sum_{i = 0}^n y_i \frac{\prod_{k \neq i} (x - x_k)}{\prod_{k \neq i} (x_i - x_k)}$$ が成り立つ。$$Q_i = \frac{y_i}{\prod_{k \neq i} (x_i - x_k)}$$ と置いておく。
各 \(i\) に対して \(Q_i\) を求めるのは \(O(n)\) で可能。
\(\prod_{k \neq i} (x - x_k)\) については、予め $$\prod_{k=0}^n (x - x_k) = x^{n+1} + c_nx^n + \ldots + c_1x + c_0$$ を求めておく。
\(c_i\) は、\(\prod_k^m (x - x_k)\) を計算して係数を求めた後、\((x - x_{m + 1})\) を掛けると係数が簡単な漸化式で表せるので、\(O(n^2)\) で求められる。
各 \(i\) について考えるときは \((x - x_i)\) で割る必要があるが、高次のほうから $$ c_n^\prime = c_{n+1} (= 1) \\ c_j^\prime = c_{j+1} + x_ic_{j + 1}^\prime$$ と計算できる。
漸化式の意味は、\(c_j^\prime\) が \(\{x_0, \ldots, x_{i - 1}, x_{i + 1}, \ldots, x_n\}\) から \(n - j\) 個選んでかけ合わせたものの総和、ということを考えると、理解できると思う。
こうして $$\prod_{k \neq i} (x - x_k) = c_n^\prime x^n + \ldots + c_1^\prime x + c_0^\prime$$ が求まったので、これに \(Q_i\) をかけたものを各 \(i\) について足し合わせればOK。

以上より、\(O(n^2)\) でこの問題が解けた。

ソースコード

#include <bits/stdc++.h>

using ll = long long;

ll modpow(ll x, ll n, const ll mod) {
    ll res = 1;
    while(n > 0) {
        if(n & 1) (res *= x) %= mod;
        (x *= x) %= mod;
        n >>= 1;
    }
    return res;
}
ll inverse(ll x, const ll mod) {
    return modpow(x, mod - 2, mod);
}

std::vector<ll> lagrange_interpolation(std::vector<ll> xs, std::vector<ll> ys, const int m) {
    const int n = xs.size();
    for(int i = 0; i < n; ++i) {
        xs[i] %= m;
        ys[i] %= m;
    }
    std::vector<ll> all_c(n + 1);
    all_c[0] = 1;
    for(int i = 0; i < n; ++i) {
        std::vector<ll> nxt(n + 1);
        for(int j = 0; j < n; ++j) {
            nxt[j + 1] = all_c[j];
        }
        for(int j = 0; j < n; ++j) {
            nxt[j] = (m + nxt[j] - xs[i] * all_c[j] % m) % m;
        }
        all_c = std::move(nxt);
    }

    std::vector<ll> c(n);
    for(int i = 0; i < n; ++i) {
        ll qi = 1;
        for(int j = 0; j < n; ++j) {
            if(i == j) continue;
            qi = qi * (m + xs[i] - xs[j]) % m;
        }
        qi = inverse(qi, m) * ys[i] % m;

        auto tmp_c = all_c;
        for(int j = n - 1; j >= 0; --j) {
            c[j] = (c[j] + qi * tmp_c[j + 1]) % m;
            tmp_c[j] = (tmp_c[j] + tmp_c[j + 1] * xs[i]) % m;
        }
    }
    return c;
}

using namespace std;

int main() {
    int p; cin >> p;
    vector<ll> xs(p), ys(p);
    for(int i = 0; i < p; ++i) {
        xs[i] = i;
        cin >> ys[i];
    }
    auto res = lagrange_interpolation(xs, ys, p);
    for(int i = 0; i < p; ++i) {
        cout << res[i] << " \n"[i + 1 == p];
    }
}

感想

係数求めるやつは初めて書いたのでちょっと混乱した。