前缀和经典问题整理

发布时间 2023-05-23 23:28:48作者: sw-lab

1、一般形式  --  区域和检索 - 数组不可变

class NumArray:

    def __init__(self, nums: List[int]):
        self.pre = [0]
        for num in nums:
            self.pre.append(self.pre[-1] + num)
        ####或者#####
        self.pre = list(accumulate(nums, initial=0))

    def sumRange(self, left: int, right: int) -> int:
        return self.pre[right + 1] - self.pre[left]

2、经典问题 --  连续数组

给定一个二进制数组 nums , 找到含有相同数量的 0 和 1 的最长连续子数组,并返回该子数组的长度。

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
        pre, m = 0, {0: -1}
        maxl = 0
        for i, num in enumerate(nums):
            pre += 1 if num == 1 else -1
            if m.get(pre, None) != None:
                maxl = max(i - m[pre], maxl)
            else:
                m[pre] = i
        return maxl

3、二维数组前缀和和差分

(1)二维数组前缀和 -- 二维区域和检索 - 矩阵不可变

 代码:

class NumMatrix:

    def __init__(self, matrix: List[List[int]]):
        self.sum_matrix = [[0] * len(matrix[0]) for _ in matrix]
        for i in range(len(matrix)):
            row_sum = 0
            for j in range(len(matrix[i])):
                row_sum += matrix[i][j]
                self.sum_matrix[i][j] = self.sum_matrix[i - 1][j] + row_sum

    def sumRegion(self, row1: int, col1: int, row2: int, col2: int) -> int:
        res = self.sum_matrix[row2][col2]
        if col1 > 0: res -= self.sum_matrix[row2][col1 - 1]
        if row1 > 0: res -= self.sum_matrix[row1 - 1][col2]
        if col1 > 0 and row1 > 0: res += self.sum_matrix[row1 - 1][col1 - 1]
        return res

(2)二维数组差分 -- 子矩阵元素加 1

class Solution:
    def rangeAddQueries(self, n: int, queries: List[List[int]]) -> List[List[int]]:
        d = [[0] * (n + 1) for _ in range(n + 1)]
        for r1, c1, r2, c2 in queries:
            d[r1][c1] += 1
            d[r2 + 1][c2 + 1] += 1
            d[r1][c2 + 1] -= 1
            d[r2 + 1][c1] -= 1
        
        ans = [[0] * (n + 1) for _ in range(n + 1)]
        for i, row in enumerate(d[:n]):
            for j, x in enumerate(row[:n]):
                ans[i + 1][j + 1] = ans[i + 1][j] + ans[i][j + 1] - ans[i][j] + x
        del ans[0]
        for row in ans:
            del row[0]
        return ans

数组差分可以看成函数微分,数组前缀和可以看成函数积分,所以差分数组的前缀和就是原数组

4、进阶问题

(1)统计回文子序列数目

class Solution:
    def countPalindromes(self, s: str) -> int:
        suf = [0] * 10
        suf2 = [0] * 100
        for d in map(int, reversed(s)):
            for j, c in enumerate(suf):
                suf2[d * 10 + j] += c
            suf[d] += 1

        ans = 0
        pre = [0] * 10
        pre2 = [0] * 100
        for d in map(int, s):
            suf[d] -= 1
            for j, c in enumerate(suf):
                suf2[d * 10 + j] -= c  # 撤销
            ans += sum(c1 * c2 for c1, c2 in zip(pre2, suf2))  # 枚举所有字符组合
            for j, c in enumerate(pre):
                pre2[d * 10 + j] += c
            pre[d] += 1
        return ans % (10 ** 9 + 7)

(2)统计上升四元组

class Solution:
    def countQuadruplets(self, nums: List[int]) -> int:
        n = len(nums)
        more = [[0] * n for _ in range(n + 1)]
        less = [[0] * n for _ in range(n + 1)]
        for j in reversed(range(n)):
            for k in reversed(range(j + 1, n)):
                if nums[j] < nums[k]:
                    more[k][j] = more[k + 1][j] + 1
                else:
                    more[k][j] = more[k + 1][j]

        for k in range(n):
            for j in range(k):
                if nums[k] > nums[j]:
                    less[j][k] = less[j - 1][k] + 1
                else:
                    less[j][k] = less[j - 1][k]
        
        res = 0
        for k in range(n):
            for j in range(k):
                if nums[k] < nums[j]:
                    res += less[j][k] * more[k][j]
        
        return res