560. Subarray Sum Equals K

LeetCode 560. Subarray Sum Equals K

Description

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals tok.

Example 1:

Input: nums = [1,1,1], k = 2
Output: 2

Example 2:

Input: nums = [1,2,3], k = 3
Output: 2

Constraints:

• 1 <= nums.length <= 2 * 10^4

• -1000 <= nums[i] <= 1000

• -10^7 <= k <= 10^7

Tags

Array, Hash Table

Solution

This is a classic prefix-sum problem. To obtain the left and right bound of a subarray whose sum is k, we can find 2 prefix-sum who meet the condition $pre[i]−pre[j−1]==k$.

In terms of implementation, we build a map m whose key is sum of first n numbers and value is the occurrences of this sum, initialized with {0: 1}. Then we iterate over nums, and:

1. calculate prefix-sum preSum;

2. update the counter of subarrays whose sum equals tok, by adding m[preSum-k] onto it;

3. add m[preSum] by 1.

Complexity

• Time complexity: $O(n)$

• Space complexity: $O(n)$

Code

func subarraySum(nums []int, k int) int {
	ans, preSum, m := 0, 0, map[int]int{0: 1}
	for _, num := range nums {
		preSum += num
		ans += m[preSum-k]
		m[preSum]++
	}
	return ans
}

Reference

1. 和为K的子数组