前缀和适用的场景是原始数组不会被修改的情况下,频繁查询某个区间的累加和。
差分数组适用的场景则是需要频繁对原始数组进行增减。
本文参考labuladong算法小抄:小而美的算法技巧:差分数组
差分数组思想
比如说,我给你输入一个数组nums
,然后又要求给区间nums[2..6]
全部加 1,再给nums[3..9]
全部减 3,再给nums[0..4]
全部加 2,再给…
一通操作猛如虎,然后问你,最后nums
数组的值是什么?
常规的思路很容易,你让我给区间nums[i..j]
加上val
,那我就一个 for 循环给它们都加上呗,还能咋样?这种思路的时间复杂度是 O(N),由于这个场景下对nums
的修改非常频繁,所以效率会很低下。
这里就需要差分数组的技巧,类似前缀和技巧构造的prefix
数组,我们先对nums
数组构造一个diff
差分数组,diff[i]
就是nums[i]
和nums[i-1]
之差:
int[] diff = new int[nums.length];
// 构造差分数组
diff[0] = nums[0];
for (int i = 1; i < nums.length; i++) {
diff[i] = nums[i] - nums[i - 1];
}
通过这个diff
差分数组是可以反推出原始数组nums
的,代码逻辑如下:
int[] res = new int[diff.length];
// 根据差分数组构造结果数组
res[0] = diff[0];
for (int i = 1; i < diff.length; i++) {
res[i] = res[i - 1] + diff[i];
}
这样构造差分数组diff
,就可以快速进行区间增减的操作,如果你想对区间nums[i..j]
的元素全部加 3,那么只需要让diff[i] += 3
,然后再让diff[j+1] -= 3
即可:
原理很简单,回想diff
数组反推nums
数组的过程,diff[i] += 3
意味着给nums[i..]
所有的元素都加了 3,然后diff[j+1] -= 3
又意味着对于nums[j+1..]
所有元素再减 3,那综合起来,是不是就是对nums[i..j]
中的所有元素都加 3 了?
只要花费 O(1) 的时间修改diff
数组,就相当于给nums
的整个区间做了修改。多次修改diff
,然后通过diff
数组反推,即可得到nums
修改后的结果。
现在我们把差分数组抽象成一个类,包含increment
方法和result
方法:
经典差分类
// 差分数组工具类
class Difference {
// 差分数组
private int[] diff;
/* 输入一个初始数组,区间操作将在这个数组上进行 */
public Difference(int[] nums) {
assert nums.length > 0;
diff = new int[nums.length];
// 根据初始数组构造差分数组
diff[0] = nums[0];
for (int i = 1; i < nums.length; i++) {
diff[i] = nums[i] - nums[i - 1];
}
}
/* 给闭区间 [i,j] 增加 val(可以是负数)*/
public void increment(int i, int j, int val) {
diff[i] += val;
if (j + 1 < diff.length) {
diff[j + 1] -= val;
}
}
/* 返回结果数组 */
public int[] result() {
int[] res = new int[diff.length];
// 根据差分数组构造结果数组
res[0] = diff[0];
for (int i = 1; i < diff.length; i++) {
res[i] = res[i - 1] + diff[i];
}
return res;
}
}
这里注意一下increment
方法中的 if 语句:
public void increment(int i, int j, int val) {
diff[i] += val;
if (j + 1 < diff.length) {
diff[j + 1] -= val;
}
}
当j+1 >= diff.length
时,说明是对nums[i]
及以后的整个数组都进行修改,那么就不需要再给diff
数组减val
了。
算法实践
力扣 370 题「区间加法」
那么我们直接复用刚才实现的Difference
类就能把这道题解决掉:
int[] getModifiedArray(int length, int[][] updates) {
// nums 初始化为全 0
int[] nums = new int[length];
// 构造差分解法
Difference df = new Difference(nums);
for (int[] update : updates) {
int i = update[0];
int j = update[1];
int val = update[2];
df.increment(i, j, val);
}
return df.result();
}
力扣 1109 航班预定统计(LeetCode)
当然,实际的算法题可能需要我们对题目进行联想和抽象,不会这么直接地让你看出来要用差分数组技巧,这里看一下力扣第 1109 题「航班预订统计」:
【python 解法】
class Solution:
def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
# 因为航班编号从1开始,故构造长为(n+1)的res和差分数组diff
res = diff = [0] * (n+1)
# 将增减变化反映到差分数组中
for book in bookings:
start, end, inc = book[0], book[1], book[2]
diff[start] += inc
# 注意判断差分右边界,若超出diff右边界则不用减去inc
if end + 1 < len(diff):
diff[end + 1] -= inc
# 将差分变化同步到res数组
for i in range(1, n+1):
res[i] = res[i-1] + diff[i]
# 注意:res[0]为0,只需返回res[1: ]
return res[1: ]
【算法思路】
函数签名如下:
int[] corpFlightBookings(int[][] bookings, int n)
这个题目就在那绕弯弯,其实它就是个差分数组的题,我给你翻译一下:
给你输入一个长度为n
的数组nums
,其中所有元素都是 0。再给你输入一个bookings
,里面是若干三元组(i,j,k)
,每个三元组的含义就是要求你给nums
数组的闭区间[i-1,j-1]
中所有元素都加上k
。请你返回最后的nums
数组是多少?
PS:因为题目说的n
是从 1 开始计数的,而数组索引从 0 开始,所以对于输入的三元组(i,j,k)
,数组区间应该对应[i-1,j-1]
。
这么一看,不就是一道标准的差分数组题嘛?我们可以直接复用刚才写的类:
int[] corpFlightBookings(int[][] bookings, int n) {
// nums 初始化为全 0
int[] nums = new int[n];
// 构造差分解法
Difference df = new Difference(nums);
for (int[] booking : bookings) {
// 注意转成数组索引要减一哦
int i = booking[0] - 1;
int j = booking[1] - 1;
int val = booking[2];
// 对区间 nums[i..j] 增加 val
df.increment(i, j, val);
}
// 返回最终的结果数组
return df.result();
}
力扣 1094 拼车(LeetCode)
【python 解法】
class Solution:
def carPooling(self, trips: List[List[int]], capacity: int) -> bool:
# 考虑到共有1001个站点,直接全量排列所有站点信息作为res和diff
res = diff = [0] * 1001
cap = capacity
for trip in trips:
start, end, num = trip[1], trip[2], trip[0]
diff[start] += num
# 注意:乘客到达end后立即下车,故此时是判断end,而不是end + 1
if end < len(diff):
diff[end] -= num
res[0] = diff[0]
for i in range(1, 1001):
res[i] = res[i-1] + diff[i]
# 任一站点处,只要有一处超载即返回False
for r in res:
if r > cap:
return False
return True
【算法思路】
函数签名如下:
boolean carPooling(int[][] trips, int capacity);
比如输入:
trips = [[2,1,5],[3,3,7]], capacity = 4
这就不能一次运完,因为trips[1]
最多只能上 2 人,否则车就会超载。
相信你已经能够联想到差分数组技巧了:trips[i]
代表着一组区间操作,旅客的上车和下车就相当于数组的区间加减;只要结果数组中的元素都小于capacity
,就说明可以不超载运输所有旅客。
但问题是,差分数组的长度(车站的个数)应该是多少呢?题目没有直接给,但给出了数据取值范围:
0 <= trips[i][1] < trips[i][2] <= 1000
车站个数最多为 1000,那么我们的差分数组长度可以直接设置为 1001:
boolean carPooling(int[][] trips, int capacity) {
// 最多有 1000 个车站
int[] nums = new int[1001];
// 构造差分解法
Difference df = new Difference(nums);
for (int[] trip : trips) {
// 乘客数量
int val = trip[0];
// 第 trip[1] 站乘客上车
int i = trip[1];
// 第 trip[2] 站乘客已经下车,
// 即乘客在车上的区间是 [trip[1], trip[2] - 1]
int j = trip[2] - 1;
// 进行区间操作
df.increment(i, j, val);
}
int[] res = df.result();
// 客车自始至终都不应该超载
for (int i = 0; i < res.length; i++) {
if (capacity < res[i]) {
return false;
}
}
return true;
}