C++ 拉格朗日插值法優化 DP

拉格朗日插值法

簡介

衆所周知，$n$ 個點 $(x_i, y_i)$（任意兩個點橫座標不相等）可以確定一個 $n - 1$ 次多項式函數 $y = f(x)$。拉格朗日插值法可以根據這 $n$ 個點求出這個多項式 $f(x)$。當然，實際應用中通常求出橫座標爲 $k$ 的點在該（ $n$ 個點確定的）多項式函數上對應的縱座標的值，代碼實現中我們也只考慮這一問題。

一個直觀的想法是利用待定係數法設 $f(x) = a_{n- 1} x ^ {n - 1} + \cdots + a_0 x ^ 0$，然後帶入 $n$ 個點得到一個 $n$ 元一次方程組，然後用 高斯消元 解得係數。但這個方法和拉格朗日插值法相比有兩個問題：一是時間複雜度爲 $O(n^3)$，而拉格朗日法時間複雜度爲 $n^2$；二就是係數可能解出小數，還可能很大，而拉格朗日法可以支持取模，並跳過“係數”這一中間步驟，直接求值。

拉格朗日插值法

我們以二次函數爲例，看一看拉格朗日插值法的具體過程：已知 $3$ 個點 $(x_1, y_1)$，$(x_2, y_2)$，$(x_3, y_3)$，求 $f(x)$。

拉格朗日（Joseph-Louis Lagrange，1736 ~ 1813）的做法非常巧妙地避開了解多元方程的過程：
令 $f_1(x)$ 表示經過點 $(x_1, 1)$，$(x_2, 0)$，$(x_3, 0)$ 的二次函數；
$f_2(x)$ 表示經過點 $(x_1, 0)$，$(x_2, 1)$，$(x_3, 0)$ 的二次函數；
$f_3(x)$ 表示經過點 $(x_0, 1)$，$(x_2, 0)$，$(x_3, 1)$ 的二次函數。
那麼 $f(x) = y_1 \cdot f_1(x) + y_2 \cdot f_2(x) + y_3 \cdot f_3(x)$。

原因很簡單，每個子函數確保經過一個點而不經過另外兩個點。

而子函數的求法很簡單，以 $f_1(x)$ 爲例：
$f_1(x) = 0$ 的兩根爲 $x = x_2$ 和 $x = x_3$，於是設 $f_1(x) = k (x - x_2) (x - x_3)$，再帶入點 $(x_1, 1)$，得到 $k = \frac{1}{(x_1 - x_2)(x_1 - x_3)}$，於是 $f_1(x) = \frac{(x - x_2) (x - x_3)}{(x_1 - x_2)(x_1 - x_3)}$。

求高次函數與求二次函數的方法同理，可得
$\begin{aligned} f_i(x) &= \prod_{1 \leq j \leq n, j \neq i} \frac{(x - x_j)}{(x_i - x_j)} \\ f(x) &= \sum_{1 \leq i \leq n} f_i(x) \end{aligned}$
於是，想求 $f(k)$ 的值，將 $k$ 代入上式即可，時間複雜度 $O(n^2)$（$n$ 爲次數）。

模板

洛谷 P4781 【模板】拉格朗日插值

#include <bits/stdc++.h>

const int MOD = 998244353;
const int MAXN = 2000;

int Mul(const int &a, const int &b) {
    return (long long)a * b % MOD;
}

int Inv(int x) {
    int y = MOD - 2, ret = 1;
    while (y) {
        if (y & 1)
            ret = Mul(ret, x);
        x = Mul(x, x);
        y >>= 1;
    }
    return ret;
}

int N, K, X[MAXN + 5], Y[MAXN + 5];

int main() {
    scanf("%d%d", &N, &K); // 求 f(K)
    for (int i = 1; i <= N; i++)
        scanf("%d%d", &X[i], &Y[i]);
    int Ans = 0;
    for (int i = 1; i <= N; i++) {
        int x = Y[i], y = 1;
        for (int j = 1; j <= N; j++)
            if (j != i) {
                x = Mul(x, (K - X[j] + MOD) % MOD);
                y = Mul(y, (X[i] - X[j] + MOD) % MOD);
            }
        Ans = (Ans + Mul(x, Inv(y))) % MOD;
    }
    printf("%d", Ans);
}

DP 優化

思路

如果沒有接觸過可能很難想到這個與 DP 的聯繫。事實上，我們可以將某一維的 DP 看作一個函數，即令$f_i(j) = dp[i][j]$（注意這個 $f_i(j)$ 與上文中的“子函數”沒有關係）那麼，如果我們要求的 $dp[i][j]$ 中的 $j$ 值很大（例如 $j = 10^9$），我們就可以只計算 $dp[i][1], dp[i][2], \cdots, dp[i][p + 1]$（$p$ 爲 $f_i(x)$ 的次數），並用點 $(1, dp[i][1])$，$(2, dp[i][2])$，…，$(p + 1, dp[i][p + 1])$ 確定多項式 $f_i(x)$，並快速求得 $f_i(j)$，即 $dp[i][j]$，時間複雜度爲 $O(p^2)$。

這類優化的難點在於要準確地計算 $p$ 的值，即 $f_i(x)$ 的次數，接下來通過例題講解如何計算 $p$。

例題一

洛谷 P4463 [集訓隊互測2012] calc

分析

發現我們只需要計算所有遞增的合法序列的值之和，然後乘上 $n!$ 即爲答案，因爲每種遞增的合法序列任意打亂順序仍然是合法的，並且原先就不同，打亂後也一定不同。

令 $dp[i][j]$ 表示：長度爲 $i$ 的所含元素值不超過 $j$ 的遞增的合法序列的值之和，考慮在第 $i$ 個位置放元素 $j$ 還是放其他小於 $j$ 的元素，本質即爲一個揹包問題，則$dp[i][j] = j \cdot dp[i - 1][j - 1] + dp[i][j - 1]$答案爲 $dp[n][k]$，然後發現 $k \leq 10^9$，不可能直接 DP。

按照上文中的方法，我們令 $f_n(i) = dp[n][i]$，所求的就是 $f_n(k)$。接下來求出多項式 $f_n(x)$ 的次數 $p$，然後我們就只需要 DP 出 $dp[n][1]$ 到 $dp[n][p + 1]$，再用拉格朗日插值法就能算出 $f_n(k)$ 了。

接下來推導 $f_n(x)$ 的次數，令 $g(n)$ 表示多項式 $f_n(x)$ 的次數：
$\begin{aligned} dp[i][j] &= j \cdot dp[i - 1][j - 1] + dp[i][j - 1] \\ f_i(j) &= j \cdot f_{i - 1}(j - 1) + f_i(j - 1) \\ f_i(j) - f_i(j - 1) &= j \cdot f_{i - 1}(j - 1) \end{aligned}$設 $f_i(x) = \sum\limits_{i = 0}^{g(n)} a_i x ^i$，將 $j$ 和 $j - 1$ 暴力代入 $f_i(j) - f_i(j - 1)$ 這個式子，發現 $a_{g(i)} j^{g(n)}$ 這個最高次項被消掉了（代入後有關最高次項的部分僅爲 $a_{g(i)} j^{g(i)} - a_{g(i)} (j - 1)^{g(i)}$）！

於是得到 $f_i(j) - f_i(j - 1)$ 的次數爲 $g(i) - 1$，又因爲 $j \cdot f_{i - 1}(j - 1)$ 的次數爲 $g(i - 1) + 1$，所以
$\begin{aligned} g(i) - 1 &= g(i - 1) + 1\\ g(i) &= g(i - 1) + 2 \end{aligned}$又因爲 $g(0) = 0$（$f_0(x) = dp[0][x] = 1$）所以 $g(n) = 2n$，證得 $f_n(x)$ 的次數爲 $2n$。

然後我們只需要用樸素的 DP 求得 $dp[n][1]$，$dp[n][2]$，…，$dp[n][2n + 1]$（注意點數要求比次數多一才能得到正確的多項式），並用拉格朗日插值法求得 $dp[n][k]$ 即可。

代碼

#include <bits/stdc++.h>

const int MAXN = 500;

int N, K, P;
int Dp[MAXN + 5][2 * MAXN + 1 + 5];

int Add(int a, const int &b) {
    a += b; return (a >= P) ? (a - P) : a;
}

int Mul(const int &a, const int &b) {
    return (long long)a * b % P;
}

int Inv(int x) {
    int y = P - 2, ret = 1;
    while (y) {
        if (y & 1)
            ret = Mul(ret, x);
        x = Mul(x, x);
        y >>= 1;
    }
    return ret;
}

int main() {
    scanf("%d%d%d", &K, &N, &P);
    int M = 2 * N + 1;
    for (int i = 0; i <= M; i++)
        Dp[0][i] = 1;
    for (int i = 1; i <= N; i++)
        for (int j = i; j <= M; j++)
            Dp[i][j] = Add(Dp[i][j - 1], Mul(Dp[i - 1][j - 1], j));
    int Ans = 0, Fac = 1;
    for (int i = 1; i <= N; i++)
        Fac = Mul(Fac, i);
    for (int i = 1; i <= M; i++) {
        int x = Dp[N][i], y = 1;
        for (int j = 1; j <= M; j++)
            if (i != j) {
                x = Mul(x, (K >= j) ? (K - j) : (K - j + P));
                y = Mul(y, (i >= j) ? (i - j) : (i - j + P));
            }
        Ans = Add(Ans, Mul(x, Inv(y)));
    }
    printf("%d", Mul(Ans, Fac));
    return 0;
}

例題二

CF995F Cowmpany Cowmpensation

題意：給定整數 $n$ 和 $D$（$1 \leq n \leq 3000$，$1 \leq D \leq 10^9$）以及一個 $n$ 個結點的樹，要求給每個結點分配一個 $[1, D]$ 之間的整數作爲權值，並且滿足父親結點權值大於等於兒子結點，求方案總數。

分析

令 $dp[u][i]$ 表示：以 $u$ 爲根的子樹中，每個結點的權值都在 $[1,i]$ 內的方案數，同樣是一個揹包
$dp[u][i] = dp[u][i - 1] + \sum_{v \text{ is a son of } u} dp[v][i - 1]$$g(n)$ 定義與上題類似，然後得到
$\begin{aligned} g(u) - 1 &= \sum_{v \text{ is a son of } u} g(v)\\ g(u) &= \sum_{v \text{ is a son of } u} g(v) + 1 \end{aligned}$注意邊界 $g(v) = [v \text{ is a leaf }]$，因爲對於一個葉子 $u$ 有 $dp[u][i] = i$。因此這就是一個子樹大小的 DP 式，於是 $g(1) = n$，暴力算得 $dp[1][1]$，$dp[1][2]$，…，$dp[1][n + 1]$，再拉格朗日即可。

代碼

#include <bits/stdc++.h>

const int MAXN = 3000;
const int MOD = 1000000007;

int N, D, M;
std::vector<int> G[MAXN + 5];

int Dp[MAXN + 5][MAXN + 5];

int Add(int a, const int &b) {
    a += b; return (a >= MOD) ? (a - MOD) : a;
}

int Mul(const int &a, const int &b) {
    return (long long)a * b % MOD;
}

int Inv(int x) {
    int y = MOD - 2, ret = 1;
    while (y) {
        if (y & 1)
            ret = Mul(ret, x);
        x = Mul(x, x);
        y >>= 1;
    }
    return ret;
}

void Dfs(int u) {
    for (int v: G[u])
        Dfs(v);
    for (int i = 1; i <= M; i++) {
        int tmp = 1;
        for (int v: G[u])
            tmp = Mul(tmp, Dp[v][i]);
        Dp[u][i] = Add(Dp[u][i - 1], tmp);
    }
}

int main() {
    scanf("%d%d", &N, &D);
    for (int i = 2; i <= N; i++) {
        int p; scanf("%d", &p);
        G[p].push_back(i);
    }
    M = N + 1;
    Dfs(1);
    int Ans = 0;
    for (int i = 1; i <= M; i++) {
        int x = Dp[1][i], y = 1;
        for (int j = 1; j <= M; j++)
            if (i != j) {
                x = Mul(x, (D >= j) ? (D - j) : (D - j + MOD));
                y = Mul(y, (i >= j) ? (i - j) : (i - j + MOD));
            }
        Ans = Add(Ans, Mul(x, Inv(y)));
    }
    printf("%d", Ans);
    return 0;
}

A trick

上面兩題的“點”的橫座標有個規律：是連續的 $p + 1$ 個正整數。結合拉格朗日插值法的分子分母的特徵，發現可以用前綴積和後綴積優化拉格朗日插值法的內層循環代碼，使時間複雜度由 $O(p^2)$ 優化爲 $O(p)$，但是複雜度的瓶頸在於開頭的樸素 DP，所以沒有提這個方法。