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{array}{rl}{f}_i(x)& =\prod _{1\le j\le n,j\ne i}\frac{(x-x_j)}{(x_i-x_j)}\\ f(x)& =\sum _{1\le i\le n}{f}_i(x)\end{array}$$
于是，想求 $f(k)$ 的值，将 $k$ 代入上式即可，时间复杂度 $O(n^2)$（$n$ 为次数）。

模板

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

#include 

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\le 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{array}{rl}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{array}$$设 $f_i(x)=\sum _{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{array}{rl}g(i)-1& =g(i-1)+1\\ g(i)& =g(i-1)+2\end{array}$$又因为 $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 

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\le n\le 3000$，$1\le D\le 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{array}{rl}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{array}$$注意边界 $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 

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，所以没有提这个方法。