Leetcode-和为k的数组(前缀和)

79 阅读2分钟

给你一个整数数组 nums 和一个整数 k ,请你统计并返回 *该数组中和为 k 的子数组的个数

暴力求解

class Solution(object):
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = 0
        n = len(nums)
        for i in range(n):
            for j in range(i,n):
                if sum(nums[i:j+1]) == k:
                    count += 1
        return count

时间复杂度为O(n^3)

暴力求解优化

其实,可以对上面方法做一些优化。因为我们求nums[i:j+1]nums[i:j+1]nums[i:j+1]的下一次nums[i:j+1+1]nums[i:j+1+1]nums[i:j+1+1],其实只比前一次多了一个nums[j+1]nums[j+1]nums[j+1]。那么我们就没必要每次都从i开始算了。

class Solution(object):
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = 0
        n = len(nums)
        for i in range(n):
            sum = 0
            for j in range(i,n):
                sum += nums[j]
                if sum == k:
                    count += 1
        return count

时间复杂度为O(n^2)

前缀和

我们通过前缀和数组保存前 n 位的和,presum[1]保存的就是 nums 数组中前 1 位的和,也就是 presum[1] = nums[0], presum[2] = nums[0] + nums[1] = presum[1] + nums[1]. 依次类推,所以我们通过前缀和数组可以轻松得到每个区间的和。

前缀和O(n^2)

class Solution(object):
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = 0
        n = len(nums)
        preSum = [0]
        tmp = 0
        for i in range(n):
            tmp += nums[i]
            preSum.append(tmp)
        
        for i in range(1, n+1):
            for j in range(i, n+1):
                if preSum[j] - preSum[i-1] == k:
                    count += 1
        return count

preSum[j] 表示的是前 j 个元素的累积和,而 preSum[j] - preSum[i-1] 表示的是索引 i-1 到 j 的子数组的和.

前缀和O(n^2)

class Solution(object):
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = 0
        n = len(nums)
        preSums = {}  # 使用普通字典
        preSums[0] = 1

        presum = 0
        for i in range(n):
            presum += nums[i]

            # 检查键是否存在,如果存在则使用对应的值,否则假设值为0
            if presum - k in preSums:
                count += preSums[presum - k]

            # 更新前缀和为 presum 的频次
            preSums[presum] = preSums.get(presum, 0) + 1

        return count
nums = [2, 4, 4, 6, 0, 3, 3]
k = 6
sol = Solution()
result = sol.subarraySum(nums, k)
print(result)

Visualization