LeetCode: 560. Subarray Sum Equals K

560. Subarray Sum Equals K

560. Subarray Sum Equals K

Summary

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

A subarray is a contiguous non-empty sequence of elements within an array.

Constraints:

  • 1 <= nums.length <= 2 * 10^4
  • -1000 <= nums[i] <= 1000
  • -10^7 <= k <= 10^7

Approach

Use a running prefix sum and store how many times each prefix sum has appeared.

If the current prefix sum is total, then any earlier prefix sum equal to total - k forms a subarray whose sum is k.

So for each position:

  1. Add the current number to total.
  2. Check how many times total - k has already appeared.
  3. Add that count to the answer.
  4. Record the new total in the map.

This works because:

  • prefix_sum[j] - prefix_sum[i] = k
  • Therefore, prefix_sum[i] = prefix_sum[j] - k

Function

Initialization:

  • mem = {0: 1}
  • total = 0
  • ans = 0

For each number n in nums:

  • total += n
  • If total - k is in mem, add mem[total - k] to ans
  • Increment mem[total]

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

  1. Read 1: total = 1, total - k = -1 is not in mem. Update mem to {0: 1, 1: 1}, ans = 0
  2. Read 1: total = 2, total - k = 0 is in mem. Add 1, so ans = 1. Update mem to {0: 1, 1: 1, 2: 1}
  3. Read 1: total = 3, total - k = 1 is in mem. Add 1, so ans = 2. Update mem to {0: 1, 1: 1, 2: 1, 3: 1}

The output is 2.

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

  1. Read 1: total = 1, total - k = 0 is in mem. Add 1, so ans = 1. Update mem to {0: 1, 1: 1}
  2. Read -1: total = 0, total - k = -1 is not in mem. Update mem to {0: 2, 1: 1}, ans = 1
  3. Read 1: total = 1, total - k = 0 is in mem twice. Add 2, so ans = 3. Update mem to {0: 2, 1: 2}

The output is 3.

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

  1. Read 1: total = 1, total - k = -2 is not in mem. Update mem to {0: 1, 1: 1}, ans = 0
  2. Read 2: total = 3, total - k = 0 is in mem. Add 1, so ans = 1. Update mem to {0: 1, 1: 1, 3: 1}
  3. Read 3: total = 6, total - k = 3 is in mem. Add 1, so ans = 2. Update mem to {0: 1, 1: 1, 3: 1, 6: 1}

The output is 2.

Complexity

Time complexity: O(N), where N is the length of the array. Space complexity: O(N) for the hash map of prefix-sum counts.