Skip to content

560. Subarray Sum Equals K

class Solution:
    """🧮 Subarray Sum Equals K.

    Problem:
    --------
    Given an integer array `nums` (values may be negative) and an integer `k`, return the
    number of non-empty contiguous subarrays whose sum equals `k`.

    Example (`nums = [1, 1, 1]`, `k = 2`, answer 2): `nums[0:2]` and `nums[1:3]`.

    Key Insight:
    ------------
    Let `P[j] = nums[0] + ... + nums[j - 1]` be the prefix sums, with `P[0] = 0` for the
    empty prefix. Every subarray sum is a difference of two prefix sums:

        sum(nums[i:j]) = P[j] - P[i]        for 0 <= i < j <= n

    So `sum(nums[i:j]) == k` iff `P[i] == P[j] - k`. Counting subarrays becomes counting
    **pairs of prefix sums**: for each end `j`, the number of subarrays ending there is
    the number of **earlier** prefixes `P[i]` (with `i < j`) equal to `P[j] - k`.

    Scan left to right with a hash map `freq[s]` = how many prefixes seen so far equal
    `s`. At each `j`, add `freq[P[j] - k]` to the answer, then record `P[j]`.

    Example (`nums = [1, 1, 1]`, `k = 2`):

          num   P[j]   P[j] - k   freq before          count
          1     1      -1         {0: 1}               0
          1     2       0         {0: 1, 1: 1}         1    (nums[0:2])
          1     3       1         {0: 1, 1: 1, 2: 1}   2    (nums[1:3])

    Why not a sliding window? Negative values break it: extending a window can lower its
    sum, so there's no rule for when to move the left end. Prefix sums work for any sign.

    Approach:
    ---------
    - `freq` starts as `{0: 1}`: the empty prefix `P[0] = 0`. Without it, subarrays that
      start at index 0 (where `P[j] == k`) are never counted.
    - Query **before** inserting the current prefix. That enforces `i < j`, so the
      subarray is non-empty. With `k = 0`, inserting first would count every `P[j]`
      against itself.
    - Store counts, not just presence: several earlier prefixes can share the same value
      (zeros or negatives), and each one starts a different subarray.

    Complexity:
    -----------
    - Time: O(n), one pass with O(1) expected hash operations
    - Space: O(n) for the prefix-sum counts
    """

    def subarraySum(self, nums: list[int], k: int) -> int:
        prefix_sum = 0
        count = 0
        # We keep track of the occurrences of each value we see
        # while computing the prefix sum
        freq = {0: 1}
        for num in nums:
            prefix_sum += num

            count += freq.get(prefix_sum - k, 0)
            freq[prefix_sum] = freq.get(prefix_sum, 0) + 1

        return count