差分数组是一种用于高效求区间和的数据结构和算法技巧,主要用于解决一个区间内的所有元素都需要增加或减少某个值的情况。
前缀和是一种在算法中经常使用的概念和技巧。它的主要思想是:对于一个序列,我们事先将其前n项的和存储起来,从而在求序列的区间和时,可以省去重复计算的时间,提高效率。
基本思想
考虑一个数组 nums,假设我们想将数组 nums 中从左边界 l 到右边界 r 之间的所有元素加上一个值 val,传统的做法是遍历这个区间,时间复杂度为 O(r-l+1)。
差分思想可以将这个操作的时间复杂度降低到 O(1)。我们构造一个与原始数组等长的差分数组 diff,其中 diff[i] 用来存储 nums[i]-nums[i-1] 的结果(第一个元素为 nums[0])。
计算过程
- 初始化差分数组 diff,所有元素都为 0
- 令 diff[l] += val, diff[r+1] -= val
- 重建原始数组:diff[i] += diff[i-1]
例题1:Leetcode1109
https://leetcode.cn/problems/corporate-flight-bookings/
这里有 n
个航班,它们分别从 1
到 n
进行编号。
有一份航班预订表 bookings
,表中第 i
条预订记录 bookings[i] = [firsti, lasti, seatsi]
意味着在从 firsti
到 lasti
(包含 firsti
和 lasti
)的 每个航班 上预订了 seatsi
个座位。
请你返回一个长度为 n
的数组 answer
,里面的元素是每个航班预定的座位总数
示例 1:
输入:bookings = [[1,2,10],[2,3,20],[2,5,25]], n = 5 输出:[10,55,45,25,25] 解释: 航班编号 1 2 3 4 5 预订记录 1 : 10 10 预订记录 2 : 20 20 预订记录 3 : 25 25 25 25 总座位数: 10 55 45 25 25 因此,answer = [10,55,45,25,25]
示例 2:
输入:bookings = [[1,2,10],[2,2,15]], n = 2 输出:[10,25] 解释: 航班编号 1 2 预订记录 1 : 10 10 预订记录 2 : 15 总座位数: 10 25 因此,answer = [10,25]
题中的first和last值即为增量的区间,seat即为增量值。并且因为first与last的取值为
1 <= firsti <= lasti <= n,初始化diff数组时len(diff) = n + 1,用于简化边界问题。
理解之后,便很容易写出代码。
class Solution(object):
def corpFlightBookings(self, bookings, n):
"""
:type bookings: List[List[int]]
:type n: int
:rtype: List[int]
"""
diff = [0] * (n+1)
for l,r,seats in bookings:
diff[l-1] += seats
diff[r] -= seats
for i in range(1,n):
diff[i] += diff[i-1]
return diff[:n]
例题2:Leetcode1094
https://leetcode.cn/problems/car-pooling/
车上最初有 capacity
个空座位。车 只能 向一个方向行驶(也就是说,不允许掉头或改变方向)
给定整数 capacity
和一个数组 trips
, trip[i] = [numPassengersi, fromi, toi]
表示第 i
次旅行有 numPassengersi
乘客,接他们和放他们的位置分别是 fromi
和 toi
。这些位置是从汽车的初始位置向东的公里数。
当且仅当你可以在所有给定的行程中接送所有乘客时,返回 true
,否则请返回 false
。
示例 1:
输入:trips = [[2,1,5],[3,3,7]], capacity = 4 输出:false
示例 2:
输入:trips = [[2,1,5],[3,3,7]], capacity = 5 输出:true
提示:
1 <= trips.length <= 1000
trips[i].length == 3
1 <= numPassengersi <= 100
0 <= fromi < toi <= 1000
1 <= capacity <= 105
增量区间为from和to,当前旅行的乘客超过capacity时,返回False。
因为to的最大取值为1000,因此差分数组的大小可初始化为1000+1或测试用例中to的最大值+1。
class Solution(object):
def carPooling(self, trips, capacity):
"""
:type trips: List[List[int]]
:type capacity: int
:rtype: bool
"""
to_max = max(ele[2] for ele in trips)
diff = [0] * (to_max+1)
for num,from_,to in trips:
diff[from_] += num
diff[to] -= num
cnt = 0
for i in range(to_max+1):
cnt += diff[i]
if cnt > capacity:return False
return True
例题3:13届蓝桥杯PythonA组 重新排序
为了达到题中的目的,可以让被查询次数多的元素为更大的元素。
通过差分数组,获取每个元素的查询次数,然后对差分数组和原数组从大到小排列,最后差分数组和原数组的元素相乘,得出新的查询结果总和。
def ques_f():
from itertools import accumulate
#使用itertools库中的accumulate函数可以快速计算前缀和
n = int(input())
A = [0] + list(map(int,input().split()))
m = int(input())
#前缀和
pres = list(accumulate(A))
#差分数组
diff = [0] * (n+2)
old,new = 0,0
for i in range(m):
l,r = map(int,input().split())
old += (pres[r] - pres[l-1])
diff[l] += 1
diff[r+1] -= 1
#前缀和
diff = list(accumulate(diff))
A.sort(reverse=True)
diff.sort(reverse=True)
for j in range(n):
new += diff[j] * A[j]
print(new-old)
ques_f()