数据结构与算法(十二) 前缀和与差分数组
大家好，我是半虹，这篇文章讲前缀和与差分数组
1 前缀和数组
前缀和数组是一个简单且巧妙的数据结构，用于快速且频繁计算数组区间和
前缀和数组通过对原始数组做预处理得到，前缀和数组中的每个元素是原始数组中前面所有元素之和
具体来说，给定原始数组 a，其长度为 n，则对应的前缀和数组为 s，其长度为 n + 1
对于前缀和数组任意位置 i，s[i] 是原始数组 a 位置 [0, i) 之和，s[0] 初始为 0
a 0 1 2 3 4 索引
1 2 3 2 1 值
+———+———+———+———+———|
+———+———+———+———| |
+———+———+———| | |
+———+———| | | |
+———| | | | |
| | | | |
s (0) 1 3 6 8 9 值
0 1 2 3 4 5 索引
建立前缀和数组的时间复杂度为 O(n)，根据前缀和数组去计算原始数组任意区间和的时间复杂度为 O(1)
举例来说，对于原始数组中区间 [i, j]，其区间和等于 s[j + 1] - s[i]
下面是一道典型的例题，leetcode303
现给定整数数组 nums，以及一系列查询，每个查询要求计算 nums 在某个索引区间的元素之和
class NumArray {
public:
vector<int> sums; // 前缀和数组
NumArray(vector<int>& nums) {
int n = nums.size();
sums.resize(n + 1);
for (int i = 1; i <= n; i++) {
sums[i] = sums[i - 1] + nums[i - 1];
}
}
int sumRange(int left , int right) {
return sums[right + 1] - sums[left];
}
};
上述思路可以进行拓展，leetcode304
若给定二维矩阵 nums，以及一系列查询，每个查询要求计算 nums 在某个子矩阵中的元素之和
给定查询
q(x1, y1, x2, y2)，其中(x1, y1)为左上角坐标，(x2, y2)为右下角坐标
上述查询可以表示为：多个以原点为起点的子矩阵的线性组合式：
q ( x 1 , y 1 , x 2 , y 2 ) = + q ( 0 , 0 , x 2 , y 2 ) − q ( 0 , 0 , x 1 − 1 , y 2 ) − q ( 0 , 0 , x 2 , y 1 − 1 ) + q ( 0 , 0 , x 1 − 1 , y 1 − 1 ) \begin{align*} q(x1,\ y1,\ x2,\ y2) = &+q(0,\ 0,\ x2,\ y2) \\ &-q(0,\ 0,\ x1 - 1,\ y2) \\ &-q(0,\ 0,\ x2,\ y1 - 1) \\ &+q(0,\ 0,\ x1 - 1,\ y1 - 1) \\ \end{align*} q(x1, y1, x2, y2)=+q(0, 0, x2, y2)−q(0, 0, x1−1, y2)−q(0, 0, x2, y1−1)+q(0, 0, x1−1, y1−1)
对于以原点为起点的子矩阵元素之和，可以通过预先计算的前缀和矩阵快速得到
设前缀和矩阵为
sums，那么有sums[i][j]等价于q(0, 0, i - 1, j - 1)
初始化
sums[0][...]和sums[...][0]为零，有sums[i][j]计算公式如下：
sums [ i ] [ j ] = + sums [ i − 1 ] [ j ] + sums [ i ] [ j − 1 ] − sums [ i − 1 ] [ j − 1 ] + nums [ i ] [ j ] \begin{align*} \text{sums}[i][j] = &+\text{sums}[i - 1][j] \\ &+\text{sums}[i][j - 1] \\ &-\text{sums}[i - 1][j - 1] \\ &+\text{nums}[i][j] \\ \end{align*} sums[i][j]=+sums[i−1][j]+sums[i][j−1]−sums[i−1][j−1]+nums[i][j]
class NumMatrix {
public:
vector<vector<int>> sums; // 前缀和矩阵
NumMatrix(vector<vector<int>>& nums) {
int m = nums.size();
int n = nums[0].size();
sums.resize(m + 1, vector<int>(n + 1));
for (int i = 1; i <= m; i++) {
for (int j = 1; j <= n; j++) {
sums[i][j] = (
+sums[i - 1][j]
+sums[i][j - 1]
-sums[i - 1][j - 1]
+nums[i - 1][j - 1]
);
}
}
}
int sumRegion(int row1, int col1, int row2, int col2) {
return (
+sums[row2 + 1][col2 + 1]
-sums[row1][col2 + 1]
-sums[row2 + 1][col1]
+sums[row1][col1]
);
}
};
下面是一些相关的题目：
- 和为
k的子数组 | leetcode560
给定一个数组 nums，以及一个整数 k，返回该数组中和为 k 的子数组数目
纯暴力
时间复杂度：
O(n^3)，其中n为数组大小
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int n = nums.size();
if (n == 0) {
return 0;
}
int ans = 0;
int sum = 0;
for (int i = 0; i < n; i++) {
for (int j = i; j < n; j++) {
// 计算 nums 在区间 [i, j] 和是否等于 k
sum = 0;
for (int m = i; m <= j; m++) {
sum += nums[m];
}
if (sum == k) {
ans += 1;
}
}
}
return ans;
}
};
前缀和
时间复杂度：
O(n^2)，其中n为数组大小
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int n = nums.size();
if (n == 0) {
return 0;
}
vector<int> sums(n + 1); // 前缀和数组
for (int i = 1; i <= n; i++) {
sums[i] = sums[i - 1] + nums[i - 1];
}
int ans = 0;
for (int i = 0; i < n; i++) {
for (int j = i; j < n; j++) {
// 计算 nums 在区间 [i, j] 和是否等于 k
if (sums[j + 1] - sums[i] == k) {
ans++;
}
}
}
return ans;
}
};
前缀和 + 哈希表
时间复杂度：
O(n^1)，其中n为数组大小
在上面的做法中，我们遍历每个数组元素将其作为区间左边界
对于每个左边界，继续往右遍历数组元素将其作为区间右边界
以此确定单区间，然后通过前缀数组快速计算区间所有元素和是否满足条件
除此之外，我们也可以遍历每个数组元素将其作为区间右边界
对于每个右边界，可以往左遍历数组元素将其作为区间左边界
以此确定单区间，然后通过前缀数组快速计算区间所有元素和是否满足条件
对于第二种做法：
假设前缀和数组为
sums，sums[i]为原始数组nums[0..i]的元素之和
假设区间右边界为
i，左边界为j，若区间满足条件，那么有sums[i] - sums[j - 1] == k
上式可以变形如下
sums[j - 1] == sums[i] - k
也就是说在统计以
i为结尾的子数组时，只需要统计有多少的j能满足sums[j] == sums[i] - k
而当我们从左往右遍历区间右边界
i时，统计j存在着很多重复子问题，例如
若右边界是
i，则要判断sums[j]是否满足等式，其中j \in [0, i - 1]
若右边界是
i + 1，还是判断sums[j]是否满足等式，其中j \in [0, i]
为此，我们可以 边遍历边计算
sums[j]，同时用哈希表存储前缀和结果及数量
另外，由于计算 前缀和只是与前一项有关，所以用变量去记录当前的前缀和即可
class Solution {
public:
int subarraySum(vector<int>& nums, int k) {
int n = nums.size();
if (n == 0) {
return 0;
}
unordered_map<int, int> map; // 前缀和及其数量
map[0]++;
int ans = 0;
int sum = 0;
for (int i = 0; i < n; i++) {
// 计算 nums 在区间 [0..i, i] 的和有多少等于 k
sum += nums[i];
if (map.find(sum - k) != map.end()) {
ans += map[sum - k];
}
map[sum]++;
}
return ans;
}
};
- 和为
k的树路径 | leetcode437
给定二叉树根 root，以及一个整数 k，返回二叉树中和为 k 的路径的数目
前缀和 + 哈希表
时间复杂度：
O(n^1)，其中n为二叉树节点数
思路与上题十分类似，只是这里需要用深度优先搜索来遍历树节点
另外在深度搜索过程，还要注意正确地维护哈希表状态
class Solution {
public:
int pathSum(TreeNode* root, int targetSum) {
unordered_map<long long int, int> map; // 前缀和及其数量
map[0]++;
return dfs(root, 0, targetSum, map);
}
int dfs(TreeNode *root, long long int sum, int targetSum, unordered_map<long long int, int>& map) { // 遍历树节点
if (!root) {
return 0;
}
int ans = 0;
sum += root -> val;
if (map.find(sum - targetSum) != map.end()) {
ans = map[sum - targetSum];
}
map[sum]++;
ans += dfs(root -> left, sum, targetSum, map);
ans += dfs(root -> right, sum, targetSum, map);
map[sum]--;
return ans;
}
};
2 差分数组
前缀和数组适用于在原始数组不被修改时，频繁地查询原始数组中某个区间的元素和
而差分数组与前缀和数组很像，主要适用于频繁地修改原始数组中某个区间的元素值
给定一个数组和一些修改操作，每次操作会修改指定区间的元素值
最后问你执行完所有修改之后，原始数组会变成什么样子，具体的例子如下：
给定数组: nums = [1, 3, 4, 7, 9]
给定修改:
1. [0, 2, +1]，表示对区间 [0, 2] 所有元素加 1
2. [1, 4, -2]，表示对区间 [1, 4] 所有元素减 2
3. ...
最后问你: nums = ...
最简单的方法莫过于暴力模拟，根据给定的每个修改操作，遍历其对应区间，并修改其中所有元素
但这样的做法时间复杂度很高，这里就需要用到差分数组来优化
与前缀和数组一样，差分数组也是对原始数组做预处理后得到的
假设原始数组为 a 且长度为 n，对应差分数组为 d 且长度为 n
那么，从原始数组构造差分数组方法如下：
先将 d[0] 初始化为 a[0]，且对于任意位置 i，有 d[i] = a[i] - a[i - 1]
a 0 1 2 3 4 索引
1 3 4 7 9 值
-———+
-———+ |
-———+ | |
-———+ | | |
| | | |
d (1) 2 1 3 2 值
0 1 2 3 4 索引
反之，从差分数组还原原始数组方法如下：
先将 a[0] 初始化为 d[0]，且对于任意位置 i，有 a[i] = a[i - 1] + d[i]（就是在算前缀和）
构建差分数组和还原原始数组的时间复杂度是 O(n)，根据差分数组去修改原始数组的时间复杂度是 O(1)
举例来说，如果想要 nums[i..j] 加 k，那么只要 diff[i] += k 且 diff[j + 1] -= k, 为什么呢
因为 diff[i] += k 意味着 nums[i..n] 加 k, 而 diff[j + 1] -= k 意味着 nums[j + 1..n] 减 k
所以综合起来就只有 nums[i..j] 加 k
下面是一些经典例题：
- 拼车 | leetcode1094
class Solution {
public:
bool carPooling(vector<vector<int>>& trips, int capacity) {
vector<int> nums(1005);
vector<int> diff(1005);
// 构造差分数组
for (vector<int> trip: trips) {
diff[trip[1]] += trip[0];
diff[trip[2]] -= trip[0];
}
// 还原原始数组
nums[0] = diff[0];
for (int i = 1; i < nums.size(); i++) {
nums[i] = nums[i - 1] + diff[i];
}
// 计算最终结果
for (int i = 0; i < nums.size(); i++) {
if (nums[i] > capacity) {
return false;
}
}
return true;
}
};
- 航班预订 | leetcode1109
class Solution {
public:
vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {
vector<int> nums(n);
vector<int> diff(n);
// 构造差分数组
for (vector<int> booking: bookings) {
diff[booking[0] - 1] += booking[2];
if (booking[1] < n) {
diff[booking[1]] -= booking[2];
}
}
// 还原原始数组
nums[0] = diff[0];
for (int i = 1; i < nums.size(); i++) {
nums[i] = nums[i - 1] + diff[i];
}
// 返回原始数组
return nums;
}
};
- 区间覆盖 | leetcode1893
class Solution {
public:
bool isCovered(vector<vector<int>>& ranges, int left, int right) {
vector<int> nums(55);
vector<int> diff(55);
// 构造差分数组
for (vector<int> range: ranges) {
diff[range[0] - 1] += 1;
diff[range[1]] -= 1;
}
// 还原原始数组
nums[0] = diff[0];
for (int i = 1; i < nums.size(); i++) {
nums[i] = nums[i - 1] + diff[i];
}
// 计算最终结果
for (int i = left - 1; i <= right - 1; i++) {
if (nums[i] == 0) {
return false;
}
}
return true;
}
};
- 区间分组 | leetcode2406
将思路转换下，所有区间的最高重叠次数，就是能将区间分开的最少组数
class Solution {
public:
int minGroups(vector<vector<int>>& intervals) {
vector<int> nums(1e6 + 5);
vector<int> diff(1e6 + 5);
// 构造差分数组
for (vector<int> interval: intervals) {
diff[interval[0] - 1] += 1;
diff[interval[1]] -= 1;
}
// 还原原始数组
nums[0] = diff[0];
for (int i = 1; i < nums.size(); i++) {
nums[i] = nums[i - 1] + diff[i];
}
// 计算最终结果
int ans = INT_MIN;
for (int i = 0; i < nums.size(); i++) {
ans = max(ans , nums[i]);
}
return ans;
}
};
对于上述的代码，我们可以做一些时间上的优化：
- 在构造差分数组时，记录下区间最小最大值，后续只需遍历这个范围
- 在还原原始数组时，同时去计算最终结果值，后续无需再次进行遍历
- 在还原原始数组时，由于只用到上一项的值，所以只用变量记录即可
class Solution {
public:
int minGroups(vector<vector<int>>& intervals) {
// 构造差分数组
// 同时记录数组的最小最大值
vector<int> diff(1e6 + 5);
int minV = INT_MAX;
int maxV = INT_MIN;
for (vector<int> interval: intervals) {
diff[interval[0] - 1] += 1;
diff[interval[1]] -= 1;
minV = min(minV, interval[0] - 1);
maxV = max(maxV, interval[1] - 1);
}
// 还原原始数组
// 同时计算结果
int sum = 0;
int ans = INT_MIN;
for (int i = minV; i <= maxV; i++) {
sum = sum + diff[i];
ans = max(ans, sum);
}
return ans;
}
};
从另一方面来说，我们也能做一些空间上的优化：
- 对差分数组的记录，无需知道所有位置的值，只需知道在哪些位置做了修改就可以，因此可用有序集合
class Solution {
public:
int minGroups(vector<vector<int>>& intervals) {
// 构造差分数组
// 差分数组使用有序集合实现
map<int, int> diff;
for (vector<int> interval: intervals) {
diff[interval[0] - 1] += 1;
diff[interval[1]] -= 1;
}
// 还原原始数组
// 同时计算结果
int sum = 0;
int ans = INT_MIN;
for (map<int, int>::iterator iter = diff.begin(); iter != diff.end(); iter++) {
sum = sum + iter -> second;
ans = max(ans, sum);
}
return ans;
}
};
- 二维差分 | leetcode2536
这是差分数组在二维情况下的应用，还是按照差分数组的思路来思考就行
只要知道在修改原始数组某区域时，差分数组怎么做相应的变化就可以了，不难发现：
若原始数组在 (x1, y1) 为左上角、(x2, y2) 为右下角的区域加 k
则差分数组在 (x1, y1) 和 (x2 + 1, y2 + 1) 加 k，在 (x1, y2 + 1) 和 (x2 + 1, y1) 减 k
这道理很简单，大家画个图就知道，只是需要明白差分数组和原始数组间的对应关系
若差分数组在 (x, y) 加 k，意味着原始数组在 (x, y) 为左上角、(m, n) 为右下角的区域 加 k
其中 m、n 分别为原始数组的宽和长
class Solution {
public:
vector<vector<int>> rangeAddQueries(int n, vector<vector<int>>& queries) {
// 构造差分数组
vector<vector<int>> diff(n, vector<int>(n));
for (vector<int> query: queries) {
diff[query[0]][query[1]] += 1;
if (query[3] + 1 < n) diff[query[0]][query[3] + 1] -= 1;
if (query[2] + 1 < n) diff[query[2] + 1][query[1]] -= 1;
if (query[2] + 1 < n && query[3] + 1 < n) diff[query[2] + 1][query[3] + 1] += 1;
}
// 按列算前缀和
for (int j = 0; j < n; j++) {
for (int i = 0; i < n; i++) {
if (i == 0) {
diff[i][j] = diff[i][j];
} else {
diff[i][j] = diff[i][j] + diff[i - 1][j];
}
}
}
// 按行算前缀和
for (int i = 0; i < n; i++) {
for (int j = 0; j < n; j++) {
if (j == 0) {
diff[i][j] = diff[i][j];
} else {
diff[i][j] = diff[i][j] + diff[i][j - 1];
}
}
}
// 返回原始数组
return diff;
}
};
好啦，本文到此结束，感谢您的阅读！
如果你觉得这篇文章有需要修改完善的地方，欢迎在评论区留下你宝贵的意见或者建议
如果你觉得这篇文章还不错的话，欢迎点赞、收藏、关注，你的支持是对我最大的鼓励 (/ω＼)