《算法竞赛进阶指南》0x03 前缀和与差分

0 阅读2分钟

前缀和与差分:区间查询与区间修改的两种基本思路(C++)

学习李煜东《算法竞赛进阶指南》0x03 时,我把这一节归纳为两个方向:前缀和预先累计信息,方便查询一段区间;差分只记录变化的边界,方便批量修改一段区间。以下例子与代码是重新整理的讲解,不是教材原文。

一维前缀和:用两个累计值求区间和

设数组从 1 开始编号,定义 S[0]=0,S[i]=A[1]+...+A[i]。每个前缀只需在上一个前缀后加一个数:

S[i]=S[i−1]+A[i].S[i]=S[i-1]+A[i].

求闭区间 [l,r] 的和时,S[r] 包含了第 1 到第 r 个数,减去 S[l-1] 就删掉了区间左侧的部分:

sum⁡(l,r)=S[r]−S[l−1].\operatorname{sum}(l,r)=S[r]-S[l-1].

例如 A=[4,1,5],前缀和依次为 S[1]=4、S[2]=5、S[3]=10。因此 [2,3] 的和是 S[3]-S[1]=10-4=6。

vector<long long> s(n + 1, 0);
for (int i = 1; i <= n; ++i) s[i] = s[i - 1] + a[i];
long long answer = s[r] - s[l - 1];

预处理是 O(n)O(n),此后每次区间求和是 O(1)O(1)。long long 用来容纳可能超过 int 范围的总和。

二维前缀和:重叠区域只算一次

令 S[i][j] 表示矩阵左上角 (1,1) 到 (i,j) 的矩形和,并把第 0 行、第 0 列设为 0。上方矩形 S[i-1][j] 与左方矩形 S[i][j-1] 相加时,左上角的 S[i-1][j-1] 被算了两次,所以需要减掉一次,再加入当前格:

S[i][j]=S[i−1][j]+S[i][j−1]−S[i−1][j−1]+A[i][j].S[i][j]=S[i-1][j]+S[i][j-1]-S[i-1][j-1]+A[i][j].

同理,左上角为 (x1,y1)、右下角为 (x2,y2) 的子矩形和为:

S[x2][y2]−S[x1−1][y2]−S[x2][y1−1]+S[x1−1][y1−1].S[x2][y2]-S[x1-1][y2]-S[x2][y1-1]+S[x1-1][y1-1].

例如矩阵 [[1,2],[3,4]] 的右下角单格,可算为 S[2][2]-S[1][2]-S[2][1]+S[1][1]=10-3-4+1=4。

教材的“激光炸弹”把目标价值累加到坐标网格,枚举边长为 R 的正方形,并用二维前缀和快速求出每个正方形内的价值。题目规定边界上的目标不被摧毁,因此建模时要把目标看作格子的中心,让格子求和与边界规则对应起来。

差分:只改区间的入口和出口

定义 B[1]=A[1],B[i]=A[i]-A[i-1]。反过来,对 B 求前缀和就还原了 A:

A[i]=∑j=1iB[j].A[i]=\sum_{j=1}^{i}B[j].

若把 A[l..r] 全部加上 d,区间内部相邻两项一起变化,差值不变;只有进入区间的 B[l] 增加 d,离开区间的 B[r+1] 减少 d:

b[l] += d;
b[r + 1] -= d;

例如 A=[2,5,5,1],差分为 B=[2,3,0,-4]。把 A[2..3] 都加 2 后,原数组成为 [2,7,7,1];差分只在第 2、4 位变化,成为 [2,5,0,-6]。这里最容易混淆的是被修改的原数组位置与差分数组的两个标记位置。

这种做法适合先完成许多区间修改,再统一求出各位置的最终值。如果修改与查询交错进行,仅靠这一份普通差分数组不够。

例题一:IncDec Sequence

每次可选一个区间整体加 1 或减 1,目标是让原数组所有数相等。等价地,差分数组中的 B[2..n] 都要变成 0;B[1] 决定最终共同的值。

设 B[2..n] 中正数的总和为 p,负数绝对值的总和为 q。一次操作可以同时抵消一个正量和一个负量,因此先配对消掉 min(p,q) 个;余下 |p-q| 个量各需一次操作:

最少操作次数=min⁡(p,q)+∣p−q∣=max⁡(p,q).\text{最少操作次数}=\min(p,q)+|p-q|=\max(p,q).

余下操作可在数组的左边或右边结束,分配方式共有 |p-q|+1 种,对应这么多种不同的最终常数数组。

用 A=[1,3,2] 检查:B[2]=2、B[3]=-1,所以 p=2,q=1。先将 [2,2] 减 1,得到 [1,2,2]。第二步可以将 [2,3] 减 1,得到 [1,1,1];也可以将 [1,1] 加 1,得到 [2,2,2]。两条方案都是 2 次操作,最终数组有 2 种。

下面的程序按通常的输入形式读取 n 和数组,并输出这两个答案:

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n;
    cin >> n;
    long long previous;
    cin >> previous;

    long long p = 0, q = 0;
    for (int i = 2; i <= n; ++i) {
        long long current;
        cin >> current;
        long long difference = current - previous;
        if (difference > 0) p += difference;
        else q -= difference;
        previous = current;
    }

    long long gap = p >= q ? p - q : q - p;
    cout << max(p, q) << '\n' << gap + 1 << '\n';
}

例题二:Tallest Cow

已知第 p 头牛最高,身高为 h。如果第 a 头与第 b 头能互相看见,那么它们之间的牛都必须比两端矮。以 h 为基准,对每一对关系,将中间的闭区间 [a+1,b-1] 的相对高度减 1。差分修改是 d[a+1]--、d[b]++,最后求前缀和并加上 h。重复出现的牛对只处理一次。

#include <algorithm>
#include <iostream>
#include <set>
#include <utility>
#include <vector>
using namespace std;

int main() {
    int n, p, h, m;
    cin >> n >> p >> h >> m;

    vector<int> d(n + 2, 0), c(n + 1, 0);
    set<pair<int, int>> seen;
    for (int k = 0; k < m; ++k) {
        int a, b;
        cin >> a >> b;
        if (a > b) swap(a, b);
        pair<int, int> relation = {a, b};
        if (seen.count(relation)) continue;
        seen.insert(relation);
        d[a + 1]--;
        d[b]++;
    }

    for (int i = 1; i <= n; ++i) {
        c[i] = c[i - 1] + d[i];
        cout << h + c[i] << '\n';
    }
}

p 被读入但没有直接参与修改:h 已经是统一的身高基准,合法关系不会要求把这头已知最高的牛降矮。使用有序 set 去重时,时间复杂度为 O(n+mlog⁡m)O(n+m\log m),空间复杂度为 O(n+m)O(n+m)。

课外练习:CF295A Greg and Array

这道题不是教材例题。它先给出 m 个区间加操作,但暂不执行;之后每个问题 [x,y] 指定要执行第 x 到第 y 个操作。起初我把“问题选操作”和“操作改数组”混在一起,后来把它们分成两个下标空间:操作编号 1..m,原数组位置 1..n。

先在操作编号上做差分:对每个问题写 cnt[x]++、cnt[y+1]--,求前缀和后,sum[i] 就是操作 i 执行的次数。再令 change = 1LL * sum[i] * d[i],表示操作 i 覆盖的每个数组元素总共要增加多少。

我先用双重循环遍历每个操作覆盖的 [l[i],r[i]],这能得到正确的结果,却在大数据上 TLE:最坏要遍历约 101010^{10} 个位置。经讲解后,我把内层循环改成第二层差分:

long long change = 1LL * sum[i] * d[i];
arrayDiff[l[i]] += change;
arrayDiff[r[i] + 1] -= change;

最后对 arrayDiff 求前缀和,加回原数组。修改后我反馈最终提交 AC;这一步优化是在讲解帮助下完成的,并非独立首解。两层差分的关键是不要把 sum[i] 误认为数组第 i 位的增量:它是第 i 个操作的执行次数。

我在这一节实际遇到的两个易错点

做前缀和的小题时,我把 A=[4,1,5] 的 S[3] 答成了 8,正确值是 10。递推式写对了,也仍需检查每一步加法。

做差分的小题时,我把 A[2..3] 加 2 误答成差分第 2、3 位都加 2。正确的是 B[2]+=2、B[4]-=2:右侧标记落在闭区间的下一位,并且符号相反。学习 IncDec Sequence 时,我也明确提出过“没看懂”;用 [1,3,2] 逐步操作后,才写出了降低 A[2] 时 B[2] 减 1、B[3] 加 1 的原因。

做题前可以先问自己:要求的是固定数组的区间和,还是许多区间修改后的最终值?前者考虑前缀和,后者考虑差分;涉及边界时,把原数组与差分数组各写一遍再核对。

参考:李煜东《算法竞赛进阶指南》,0x03“前缀和与差分”,书中第 21~25 页。