[leetcode] 560. Subarray Sum Equals K @ python

原题

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

Example 1:
Input:nums = [1,1,1], k = 2
Output: 2
Note:
The length of the array is in range [1, 20,000].
The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

解法

参考: 花花酱 LeetCode 560. Subarray Sum Equals K
构造字典count, key是从index 0 到index i 的和, value是这个和出现的次数, 假设当前值是cur, 那么我们只需寻找cur-k在字典里的次数, 因为如果从index 0 到index j 的和为cur - k, 那么从j+1到 i 的subarray的和一定为k. 初始化count[0]为1的原因是假设符合题意subarray是从index 0 开始, 那么它之前的和(为0)的次数需要设置为1, 否则这个subarray无法累加到结果里.

Time: O(n)
Space: O(1)

代码

class Solution:
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = {0:1}
        cur, res = 0, 0
        for n in nums:
            cur += n
            res += count.get(cur - k, 0)
            count[cur] = count.get(cur, 0) + 1
            
        return res