前缀和与差分 | 数组区间查询的利器
前缀和与差分 | 数组区间查询的利器
引言
前缀和(Prefix Sum)与差分(Difference Array)是数组处理中两种重要且互补的技术。前缀和用于快速计算数组区间元素的和,而差分用于快速对数组区间进行相同的加减操作。这两种技术看似简单,却在LeetCode和实际工程中有着广泛的应用。
前缀和的核心思想是将数组的累计信息预先计算并存储,使得后续的区间查询操作可以在 O(1) 时间内完成。差分则是前缀和的"逆操作",它将区间更新转换为单点更新,大幅提高了批量区间更新的效率。本文将系统讲解前缀和与差分的原理、实现和应用场景。
前缀和基础
一维前缀和
一维前缀和是最基础的形式。给定数组 nums,其前缀和数组 prefix 满足:prefix[i] = nums[0] + nums[1] + ... + nums[i]。通常我们定义 prefix[0] = 0,这样 prefix[i] 表示 nums[0...i-1] 的和,即前 i 个元素的和。
前缀和数组的构造只需 O(n) 时间:prefix[0] = 0,然后对于 i 从 1 到 n,prefix[i] = prefix[i-1] + nums[i-1]。
前缀和的核心应用
前缀和的最大价值在于快速计算任意区间的和。对于区间 [l, r](包含两端)的元素和,可以直接通过 prefix[r+1] - prefix[l] 计算。这是因为 prefix[r+1] = nums[0] + ... + nums[r],prefix[l] = nums[0] + ... + nums[l-1],两者相减正好得到 nums[l] + ... + nums[r]。
这个公式的时间复杂度是 O(1),而如果直接遍历计算区间和需要 O(n) 时间。当需要大量区间查询时,前缀和可以将总时间复杂度从 O(n*q) 降低到 O(n+q),其中 q 是查询次数。
代码实现
class PrefixSum: def __init__(self, nums): self.prefix = [0] * (len(nums) + 1) for i in range(len(nums)): self.prefix[i + 1] = self.prefix[i] + nums[i] def query(self, left, right): return self.prefix[right + 1] - self.prefix[left]差分数组
差分的定义
给定数组 nums,其差分数组 diff 满足:diff[i] = nums[i] - nums[i-1](对于 i > 0),diff[0] = nums[0]。更实用的定义是:当我们想对区间 [l, r] 的所有元素加上 val 时,只需要执行 diff[l] += val,diff[r+1] -= val。然后通过对 diff 求前缀和,就可以得到更新后的数组。
差分的核心价值在于:将区间更新操作从 O(n) 降低到 O(1)。当我们需要对数组进行大量区间更新时,差分数组可以将总时间复杂度从 O(n*u) 降低到 O(n+u),其中 u 是更新次数。
差分与前缀和的关系
差分和前缀和是一对互逆的操作。对数组求前缀和得到前缀和数组,对前缀和数组求差分得到原数组。反过来,对数组求差分得到差分数组,对差分数组求前缀和得到原数组。
这种互逆关系使得差分数组特别适合处理"先批量更新,后查询最终结果"的场景。我们可以先将所有更新操作记录在差分数组中(每个操作只需 O(1)),最后一次性计算前缀和得到最终结果。
代码实现
class DifferenceArray: def __init__(self, nums): self.diff = [0] * (len(nums) + 1) if len(nums) > 0: self.diff[0] = nums[0] for i in range(1, len(nums)): self.diff[i] = nums[i] - nums[i - 1] def update(self, left, right, val): self.diff[left] += val if right + 1 < len(self.diff): self.diff[right + 1] -= val def get_result(self): result = [0] * (len(self.diff) - 1) result[0] = self.diff[0] for i in range(1, len(self.diff) - 1): result[i] = result[i - 1] + self.diff[i] return result二维前缀和
二维前缀和的定义
二维前缀和是前缀和概念在二维数组上的扩展。给定二维数组 matrix,其二维前缀和 prefix[i][j] 表示以 (0, 0) 为左上角,(i, j) 为右下角的矩形区域中所有元素的和。
二维前缀和的计算公式为:prefix[i][j] = matrix[i][j] + prefix[i-1][j] + prefix[i][j-1] - prefix[i-1][j-1]。这个公式的推导基于容斥原理:当前元素加上上方矩形和左方矩形,但需要减去重复加的左上角矩形。
二维区间查询
使用二维前缀和,可以 O(1) 时间计算任意子矩阵的和。对于子矩阵 [(row1, col1), (row2, col2)],其和为:prefix[row2][col2] - prefix[row1-1][col2] - prefix[row2][col1-1] + prefix[row1-1][col1-1]。
代码实现
def build_2d_prefix(matrix): if not matrix or not matrix[0]: return [[]] m, n = len(matrix), len(matrix[0]) prefix = [[0] * (n + 1) for _ in range(m + 1)] for i in range(1, m + 1): for j in range(1, n + 1): prefix[i][j] = (matrix[i - 1][j - 1] + prefix[i - 1][j] + prefix[i][j - 1] - prefix[i - 1][j - 1]) return prefix def query_2d(prefix, row1, col1, row2, col2): return (prefix[row2 + 1][col2 + 1] - prefix[row1][col2 + 1] - prefix[row2 + 1][col1] + prefix[row1][col1])LeetCode 303:区域和检索
题目描述
LeetCode 303 要求我们设计一个类 NumArray,支持 sumRange(left, right) 方法,返回数组 nums 从索引 left 到 right(含)的元素和。题目要求多次调用 sumRange 方法,因此使用前缀和优化是必要的。
解决方案
class NumArray: def __init__(self, nums): self.prefix = [0] for num in nums: self.prefix.append(self.prefix[-1] + num) def sumRange(self, left, right): return self.prefix[right + 1] - self.prefix[left]这个解决方案的时间复杂度:初始化 O(n),每次查询 O(1)。空间复杂度 O(n)。相比不使用前缀和的暴力方法(每次查询 O(n)),总时间复杂度从 O(n*q) 降低到 O(n+q)。
LeetCode 304:二维区域和检索
题目描述
LeetCode 304 是 303 的二维版本,要求在一个二维矩阵中检索子矩阵的元素和。同样地,我们需要使用二维前缀和来优化多次查询。
解决方案
class NumMatrix: def __init__(self, matrix): if not matrix or not matrix[0]: self.prefix = [[]] return m, n = len(matrix), len(matrix[0]) self.prefix = [[0] * (n + 1) for _ in range(m + 1)] for i in range(1, m + 1): for j in range(1, n + 1): self.prefix[i][j] = (matrix[i - 1][j - 1] + self.prefix[i - 1][j] + self.prefix[i][j - 1] - self.prefix[i - 1][j - 1]) def sumRegion(self, row1, col1, row2, col2): return (self.prefix[row2 + 1][col2 + 1] - self.prefix[row1][col2 + 1] - self.prefix[row2 + 1][col1] + self.prefix[row1][col1])LeetCode 1109:航班预订统计
题目描述
LeetCode 1109 模拟航班预订系统。有 n 个航班,编号从 1 到 n。一共有 m 条预订记录,每条记录是 [first, last, seats],表示从 first 到 last 站(包括两端)的航班预订了 seats 个座位。需要返回每个航班的最终预订座位数。
差分应用
这道题是差分数组的经典应用。每条预订记录 [first, last, seats] 对应于对区间 [first, last] 加 seats。我们可以使用差分数组记录这些更新,最后通过求前缀和得到最终结果。
def corpFlightBookings(bookings, n): diff = [0] * (n + 1) for first, last, seats in bookings: diff[first - 1] += seats diff[last] -= seats result = [0] * n result[0] = diff[0] for i in range(1, n): result[i] = result[i - 1] + diff[i] return result注意数组索引从 0 开始,而航班编号从 1 开始,因此需要将 first 和 last 减 1 转换为数组索引。
LeetCode 1094:拼车
题目描述
LeetCode 1094 模拟拼车问题。车上最初有 capacity 个空座位。一个行程数组 trips 表示乘客的上车和下车地点,第 i 个行程是 [num_passengers, start, end],表示有 num_passengers 名乘客要在 start 站上车,在 end 站下车(不包含 end 站下车)。问车上是否会有超过 capacity 名乘客。
差分解法
def carPooling(trips, capacity): diff = [0] * 1001 for num, start, end in trips: diff[start] += num diff[end] -= num current = 0 for i in range(1001): current += diff[i] if current > capacity: return False return True使用差分数组记录每个站点的乘客变化,然后通过前缀和计算每站的乘客数,判断是否超过 capacity。
前缀和的高级应用
哈希表结合
前缀和可以与哈希表结合解决更复杂的问题。在 LeetCode 560(和为 K 的子数组)中,我们需要找出和等于 K 的连续子数组的数量。可以使用前缀和加哈希表的方法:遍历数组,计算当前前缀和,然后在哈希表中查找有多少个前缀和等于 current_sum - K。
def subarraySum(nums, k): prefix_sum = 0 count = 0 prefix_map = {0: 1} for num in nums: prefix_sum += num count += prefix_map.get(prefix_sum - k, 0) prefix_map[prefix_sum] = prefix_map.get(prefix_sum, 0) + 1 return count前缀和与动态规划
前缀和实际上是动态规划的一种特例形式。在很多动态规划问题中,前缀和状态转移方程的基础。理解前缀和有助于理解更复杂的动态规划问题。
复杂度分析
时间复杂度
前缀和构造:O(n)
差分数组构造:O(n)
单次区间查询:O(1)
单次区间更新(使用差分):O(1)
还原差分数组:O(n)
二维前缀和构造:O(mn)
空间复杂度
一维前缀和:O(n)
一维差分:O(n)
二维前缀和:O(mn)
总结
前缀和与差分是处理数组区间查询和更新的两种核心技术。前缀和将区间查询从 O(n) 优化到 O(1),差分将区间更新从 O(n) 优化到 O(1)。两者的结合使用可以高效解决"多次区间更新后查询"或"多次区间查询"的问题。
在实际应用中,前缀和常用于数据统计、图像处理、声音信号处理等领域。差分则常用于需要批量更新后一次性查询结果的场景,如航班预订、考勤统计等。掌握这两种技术对于算法能力的提升和解决实际问题都有重要意义。
