560. Subarray Sum Equals K (contains negative numbers)

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:

1. The length of the array is in range [1, 20,000].
2. The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

Thoughts:

1. Sum = k -> use hash table (hash map)
2. Find subarray #: record preSum , # if occurence with that preSum val
3. Initialization: (0, 1)

Code

class Solution {
    public int subarraySum(int[] nums, int k) {
        int sum = 0, result = 0;
        Map<Integer, Integer> preS = new HashMap<Integer, Integer>();
        preS.put(0,1);
        for(int num: nums){
            sum += num;
            if(preS.containsKey(sum-k)){
                result += preS.get(sum - k); // find how many preSum has value sum - k so that the right part sum is k.
            }

            // put prefix sum for the future
            preS.put(sum, preS.getOrDefault(sum,0) + 1);

        }

       return result; 
    }
}

Python

class Solution(object):
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        count = collections.Counter()
        count[0] = 1
        ans = preS = 0
        for num in nums:
            preS += num
            ans += count[preS - k]
            count[preS] +=1
        return ans