【算法-面试】差分数组和前缀和专题

1. 差分数组

# coding = "utf-8"

'''
差分数组
差分数组的主要适用场景是频繁对原始数组的某个区间的元素进行增减
'''


class Diff:
    '''
    思路：
        1.
        2.
        3.
    '''

    def __init__(self, nums):
        self.nums = nums
        self.n = len(nums)
        self.__diff_nums = [0 for _ in range(self.n)]

    def diff(self):
        self.__diff_nums[0] = self.nums[0]
        for i in range(1, self.n):
            self.__diff_nums[i] = self.nums[i] - self.nums[i - 1]
        print('diff: ', self.__diff_nums)

    def incr(self, i, j, v):
        '''
        在[i...j]之间进行增加v操作，v可以为负数
        '''
        # [i...n]之间的都要加v，第i个差分需要加v
        self.__diff_nums[i] += v
        if j + 1 < self.n:
            # [j+1...n]之间的都要减去v，因此第j+1个差分需要-v
            self.__diff_nums[j + 1] -= v
        print('incr: ', self.__diff_nums)

    def result(self):
        res = [0 for _ in range(self.n)]

        res[0] = self.__diff_nums[0]
        for i in range(self.n):
            res[i] = res[i - 1] + self.__diff_nums[i]

        return res


def getModifiedArray(length, updates):
    '''
    假设你有⼀个⻓度为 n 的数组，初始情况下所有的数字均为 0，你将会被给出 k 个更新的操作。
    其中，每个操作会被表示为⼀个三元组：[startIndex, endIndex, inc]，你需要将⼦数组 A[startIndex ... endIndex]（包括 startIndex 和 endIndex）增加 inc。
    请你返回 k 次操作后的数组
    leetcode: 370. 区间加法
    input:length = 5, updates = [[1,3,2],[2,4,3],[0,2,-2]]
    output:[-2,0,3,5,3]
    思路：
        1.
        2.
        3.
    '''
    # nums = [0 for _ in range(length)]
    nums = [1, 2, 3, 4, 5]
    d = Diff(nums)
    d.diff()
    for i in range(len(updates)):
        d.incr(updates[i][0], updates[i][1], updates[i][2])
    print(d.result())
    return d.result()


def carPooling(trips, capacity):
    '''
    你是⼀个开公交⻋的司机，公交⻋的最⼤载客量为 capacity，沿途要经过若⼲⻋站，
    给你⼀份乘客⾏程表int[][] trips，其中 trips[i] = [num, start, end] 代表着有 num 个旅客要从站点 start 上⻋，
    到站点 end 下⻋，请你计算是否能够⼀次把所有旅客运送完毕（不能超过最⼤载客量 capacity）。
    车站的个数: 0 <= trips[i][1] < trips[i][2] <= 1000
    leetcode: 1094. 拼⻋
    input: trips = [[2,1,5],[3,3,7]], capacity = 4
    output: false
    思路：
        1.
        2.
        3.
    '''
    nums = [0 for _ in range(1001)]
    df = Diff(nums)

    for t in trips:
        # 注意end是下车点，也就是end的时候，不包含下车乘客的数量，因此需要end-1，即t[2]-1
        df.incr(t[1], t[2] - 1, t[0])

    res = df.result()
    print(res)
    for r in res:
        if r > capacity:
            return False

    return True


def corpFlightBookings(bookings, n):
    '''
    这⾥有 n 个航班，它们分别从 1 到 n 进⾏编号。有⼀份航班预订表 bookings，
    表中第 i 条预订记录 bookings[i] = [firsti, lasti, seatsi] 意味着在从 firsti 到 lasti（包含 firsti 和 lasti）的每个航班上预订了 seatsi 个座位。
    请你返回⼀个⻓度为 n 的数组 answer，⾥⾯的元素是每个航班预定的座位总数。
    leetcode: 1109. 航班预订统计
    input: bookings = [[1,2,10],[2,3,20],[2,5,25]], n = 5
    output: [10,55,45,25,25]
    思路：
        1. 利用差分数组
        2. 注意是从1...n编号
        3.
    '''
    nums = [0 for _ in range(n)]
    df = Diff(nums)
    for b in bookings:
        df.incr(b[0] - 1, b[1] - 1, b[2])

    print(df.result())
    return df.result()


if __name__ == "__main__":
    getModifiedArray(5, [[1, 3, 2], [2, 4, 3], [0, 2, -2]])
    corpFlightBookings([[1, 2, 10], [2, 3, 20], [2, 5, 25]], 5)
    carPooling([[2, 1, 5], [3, 3, 7]], 4)

2. 前缀和

# coding = "utf-8"

'''
前缀和
'''


class NumArray:
    '''
    给你输⼊⼀个整数数组 nums，请你实现 NumArray 类：
        1、NumArray(int[] nums) 使⽤数组 nums 初始化对象
        2、int sumRange(int i, int j) 返回数组 nums 从索引 i 到 j（i ≤ j）范围内元素的总和，包含 i，j 两点（也就是 sum(nums[i], nums[i + 1], ... , nums[j])）
    leetcode: 303. 区域和检索 - 数组不可变
    input:  ["NumArray", "sumRange", "sumRange", "sumRange"]
            [[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]]
    output: [null, 1, -1, -3]
    思路：
        1. 用空间换时间
        2. 先求出前缀和，存入一个新的数组中，sum[i]就是nums[0...i-1]的和
        3. [i...j]之间的和就是sum(nums[0...j])-sum(nums[0...i-1])=sum[j+1]-sum[i]
        4. 这样做的目的是为了减少调用sumRange时，反复使用for循环带来的开销
        5. 注意：
    '''

    def __init__(self, nums):
        self.nums = nums
        self.pre_sum = [0] * len(nums)
        self.pre_sum[0] = self.nums[0]
        for i in range(1, len(nums)):
            self.pre_sum[i] = self.pre_sum[i - 1] + self.nums[i - 1]
        print(self.pre_sum)

    def sumRange(self, i, j):
        return self.pre_sum[j + 1] - self.pre_sum[i]


class NumMatrix:
    '''
    给定⼀个⼆维矩阵 matrix，其中的⼀个⼦矩阵⽤其左上⻆坐标 (row1, col1) 和右下⻆坐标(row2, col2) 来表示。
    请你实现 NumMatrix 类：
        1、NumMatrix(int[][] matrix) 给定整数矩阵 matrix 进⾏初始化
        2、int sumRegion(int row1, int col1, int row2, int col2) 返回左上⻆ (row1, col1) ，右下⻆ (row2, col2) 所描述的⼦矩阵的元素总和。
    leetcode: 304. ⼆维区域和检索 - 矩阵不可变
    input:  ["NumMatrix","sumRegion","sumRegion","sumRegion"]
            [[[[3,0,1,4,2],[5,6,3,2,1],[1,2,0,1,5],[4,1,0,1,7],[1,0,3,0,5]]],
            [2,1,4,3],[1,1,2,2],[1,2,2,4]]
    output: [null, 8, 11, 12]
    思路：
        1. pre_sum[i][j]代表矩阵[0,0,i-1,j-1]的和
        2. pre_sum[0][.]和pre_sum[.][0]都为0，相当于行和列方向上垫高了一层
        3. 所求面积就是大矩阵减去两个小矩阵再加上多减去的重复矩阵
    '''

    def __init__(self, matrix):
        m, n = len(matrix), len(matrix[0])
        self.print_matrix(matrix)
        self.pre_sum = [[0 for _ in range(n + 1)] for _ in range(m + 1)]
        for i in range(1, m + 1):
            for j in range(1, n + 1):
                self.pre_sum[i][j] = self.pre_sum[i - 1][j] + self.pre_sum[i][j - 1] + matrix[i - 1][j - 1] - \
                                     self.pre_sum[i - 1][j - 1]
        self.print_matrix(self.pre_sum)

    def print_matrix(self, m):
        for i in range(len(m)):
            print(i, m[i])
        print('\n')

    def sumRegion(self, x1, y1, x2, y2):
        return self.pre_sum[x2 + 1][y2 + 1] - self.pre_sum[x1][y2 + 1] - self.pre_sum[x2 + 1][y1] + self.pre_sum[x1][y1]


def subarraySum(nums, k):
    '''
    给你⼀个整数数组 nums 和⼀个整数 k，请你统计并返回该数组中和为 k 的连续⼦数组的个数。
    leetcode: 560 和为 K 的子数组
    input:nums = [1,1,1], k = 2
    output:2
    思路：
        1.
        2.
        3.
    '''
    n = len(nums)
    # map {前缀和:该前缀和出现的次数}
    pre_sum = {0: 1}
    res, sum0_i = 0, 0
    for i in range(n):
        sum0_i += nums[i]
        # 这是我们想找的前缀和 nums[0..j]
        sum0_j = sum0_i - k

        # 如果前面有这个前缀和，则直接更新答案
        if sum0_j in pre_sum:
            # 每一个i对应一种新的情况，因此需要加上hash表中的值，而不是+1
            res += pre_sum.get(sum0_j)

        # 更新hash表的值
        if sum0_i in pre_sum:
            pre_sum[sum0_i] = pre_sum.get(sum0_i) + 1
        else:
            pre_sum[sum0_i] = 1
    print(res)
    return res


if __name__ == "__main__":
    # n = NumArray(nums=[-2, 0, 3, -5, 2, -1])
    # print(n.sumRange(0, 2))
    # nm = NumMatrix([[3, 0, 1, 4, 2], [5, 6, 3, 2, 1], [1, 2, 0, 1, 5], [4, 1, 0, 1, 7], [1, 0, 3, 0, 5]])
    # subarraySum([1, 1, 1], 2)
    subarraySum([1, 2, 3, -1, 1,-1,1], 5)