基础算法10 - 前缀和

冷淡der熊
1,075 阅读5分钟

应用

  • 当题目要求子数组和 或者 连续的子数组，可以考虑用prefix-sum
  • sum(i~j) = preSum[j + 1] - preSum[i]（当preSum的首位为0）
  • 常被应用于：
    • 2-sum系列（和/差/余数/0）
    • range sum
    • 滑动窗口
    • 单调队列
    • 扫描线（少量）
    • rolling hash（略）

2-Sum系列

2-sum 变体：和 / 差 / 余数 / 0

560. 和为 K 的子数组（Medium）

image.png

Solu：

  • 最优解O(N)：
    • 2-sum的变种：计算出preSum[i]后，找两数之差preSum[i] - preSum[j] == k (j < i)

Code：

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        preSum = 0
        dic = {0: 1}  # key: sum, value: frequency
        res = 0
        for num in nums:
            preSum += num
            res += dic.get(preSum - k, 0)
            dic[preSum] = dic.get(preSum, 0) + 1
        return res

974. 和可被 K 整除的子数组（Medium）

image.png

Solu：

  • 当相同的余数再次出现时，两次出现相同余数的位置i和j之间的subarray必定是可以背整除的

Code：

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        preSum, mod = 0, 0
        count = 0
        d = {0: 1}
        for num in nums:
            preSum += num
            mod = preSum % k
            count += d.get(mod, 0)  # 相同的余数再次出现
            d[mod] = d.get(mod, 0) + 1
        return count

523. 连续的子数组和（Medium）

image.png

Solu：

  • modDict只保留每个mod第一次出现的index
  • modDict初始化为{0 : -1}：

    • index = -1：避免nums[0] = m * k就直接返回true了（invalid case）

Code：

class Solution:
    def checkSubarraySum(self, nums: List[int], k: int) -> bool:
        preSum, mod = 0, 0
        modDic = {0: -1}  # key: mod, value: preSum
        for i, num in enumerate(nums):
            preSum += num
            mod = preSum % k
            if mod in modDic:
                if i - modDic[mod] > 1:
                    return True
            else:  # 只保留第一次出现的位置
                modDic[mod] = i
        return False

525. 连续数组（Medium）

image.png

Solu：

  • 重点：把0翻转成-1
    • if 在subarray中满足#0 == #-1 == #1，then 这个subarray的sum == 0
  • dict只保留每个preSum第一次出现的index
    • dict初始化为{0 : -1}：为了不错过第一次出现#0 == #1的subarray

Code：

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
        preSum = 0
        dic = {0: -1}  # key: sum, value: index
        ans = -1
        for i, num in enumerate(nums):  # 把0当成-1，1还是1，方便计算
            preSum += (1 if num == 1 else -1)
            if preSum in dic:
                ans = max(ans, i - dic[preSum])
            else:  # 只保留preSum第一次出现的位置
                dic[preSum] = i
        return ans

Range Sum

「差分数组」 &「Lazy Propagation」（延迟更新）：

  • 对于一个数组A[0...n-1]，其差分数组B[0...n-1]定义为：
    • 当i = 0时, B[i] = A[i]
    • 当i > 0时, B[i] = A[i] - A[i - 1]
  • 求差分数组的过程，就是求前缀和数组的逆过程
    • 恢复为原数组时，A[i] = B[0] + ... + B[i]
  • 对原数组A做更新操作[L, R, val]，反映到差分数组B上，只用更新区间两端B[L]和B[R+1]的值：B[L]' = B[L] + val and B[R + 1]' = B[R + 1] - val

    • if i = L, then B[i]' = A[i]' - A[i - 1]' = (A[i] + val) - A[i - 1] = B[i] + val

    • if L < i ≤ R, then B[i]' = A[i]' - A[i - 1]' = (A[i] + val) - (A[i - 1] + val) = B[i]

      • 对中间部分的更新操作，由于相邻两项之间相减（求差分）而互相抵消

    • if i = R + 1, then B[i]' = A[i]' - A[i - 1]' = A[i] - (A[i - 1] + val) = B[i] - val

1D Array

370. 区间加法（Medium）

image.png

Solu：

  • 差分数组 + lazy propogation，略

image.png

Code：

class Solution:
    def getModifiedArray(self, length: int, updates: List[List[int]]) -> List[int]:
        res = [0] * length
        for l, r, val in updates:  # 更新差分数组
            res[l] += val
            if r + 1 < length:
                res[r + 1] -= val
        # 求差分数组的前缀和 = 恢复为（更新后的）原数组
        for i in range(1, length):
            res[i] += res[i - 1]
        return res

2D Marrix

304. 二维区域和检索 - 矩阵不可变（Medium）

image.png

Solu：

image.png

对于原矩阵A来说，其2D差分数组B可定义为：

  • B[i][j] = sum(A[0][0] ~ A[i][j]构成的矩形)
    • B[i][j] = B[i-1][j] + B[i][j-1] + A[i][j] - B[i-1][j-1]
    • sumRegion((a,b)->(c,d)) = B[c][d] - B[c][b-1] - B[a-1][d] + B[a-1][b-1]

PS：可以补dummy的0，避免讨论边界情况

Code：

class NumMatrix:
    
    def __init__(self, matrix: List[List[int]]):
        self.sumMatrix = [[0] * (len(matrix[0]) + 1) for _ in range(len(matrix) + 1)]
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                self.sumMatrix[i + 1][j + 1] = self.sumMatrix[i + 1][j] + self.sumMatrix[i][j + 1] - self.sumMatrix[i][
                    j] + matrix[i][j]
    
    def sumRegion(self, row1: int, col1: int, row2: int, col2: int) -> int:
        return self.sumMatrix[row2 + 1][col2 + 1] - self.sumMatrix[row2 + 1][col1] - self.sumMatrix[row1][col2 + 1] + \
               self.sumMatrix[row1][col1]

滑动窗口

滑动窗口的“左缩进”是不论subarray sum大小直接缩进：对于全正数是ok的，因为缩进一定会让sum减小

209. 长度最小的子数组（Medium）

image.png

Solu：

  • 本题核心是『滑动窗口』；preSum作为一个rolling sum表示当前窗口内nums[L:cur_idx]内部数字的和

Code：

code 1: 一个变量代表当前var preSum = sum(nums[l:i+1])

class Solution:
    def minSubArrayLen(self, target: int, nums: List[int]) -> int:
        preSum, l = 0, 0
        minLen = len(nums) + 1
        for i, num in enumerate(nums):
            preSum += num
            while preSum >= target:  # 此时preSum = sum(nums[l:i+1])
                minLen = min(minLen, i - l + 1)
                preSum -= nums[l]
                l += 1
        return minLen if minLen <= len(nums) else 0

code 2: 修改原数组，成为preSum Array

class Solution:
    def minSubArrayLen(self, target: int, nums: List[int]) -> int:
        nums.insert(0, 0)
        left = 0
        ans = len(nums) + 1
        for i in range(1, len(nums)):
            nums[i] += nums[i - 1]
            if nums[i] - nums[left] >= target:
                while nums[i] - nums[left] >= target:
                    left += 1
                ans = min(ans, i - left + 1)
        return ans if ans != len(nums) + 1 else 0

单调队列

sliding widow还是monqueue，取决于nums中是否有负数

862. 和至少为 K 的最短子数组（Hard）❤️

image.png

Solu：

  • 因为K > 0，所以必定是一个较大的preSum - 一个较小的preSum -> 维护一个monoqueue保证window内sum值的单调性
  • 正确性：不用担心被pop掉的index可能会和未来的某个index'构成符合条件的subarray，因为monoque的单调递增性质，被保留下来的index''更可以和未来的某个index构成符合条件的subarray，且此时len(subarray[index'':index'])更小
    • if preSum[i'] - preSum[i] >= k && preSum[i''] < preSum[i], then preSum[i'] - preSum[i''] >= k必然成立

Code：

class Solution:
    def shortestSubarray(self, nums: List[int], k: int) -> int:
        nums = [0] + nums
        for i in range(1, len(nums)):
            nums[i] += nums[i - 1]  # 构建preSum数组
        res = len(nums) + 1
        deq = collections.deque()  # 双端队列存储前缀和的index
        for i, num in enumerate(nums):
            while deq and nums[deq[-1]] >= num:  # 右出queue，保持monoqueue的单调递增性
                deq.pop()
            while deq and num - nums[deq[0]] >= k:  # 左出queue，不断缩小window size
                res = min(res, i - deq.popleft())
            deq.append(i)  # 把index放入monoqueue
        return res if res <= len(nums) else -1

最大子序和/子序积

⚠️ 重要结论：最大子序列的和/积使用DP来做的！！而不是前缀和/积！！

  • 最长/短子序列：滑动窗口（+前缀和）
  • 最大子序列：DP

152. 乘积最大子数组（Medium）

image.png

Solu：DP

对于乘法，我们需要注意：负数乘以负数，会变成正数。

  • 因此，我们需要维护两个变量：当前的最大值（以nums[i]结尾） 以及 当前的最小值（以nums[i]结尾）
    • 最小值可能为负数，但没准下一步乘以一个负数，当前的最大值就变成最小值，而最小值则变成最大值了

Code：一维滚动数组优化

class Solution:
    def maxProduct(self, nums: List[int]) -> int:
        maxP, minP = nums[0], nums[0]
        dp = nums[0]
        for i in range(1, len(nums)):
            tmp = maxP
            maxP = max(maxP * nums[i], minP * nums[i], nums[i])
            minP = min(tmp * nums[i], minP * nums[i], nums[i])
            dp = max(dp, maxP)
        return dp

Reference：

avatar
文章
阅读
粉丝