Subarray Sum Equals K [Medium] — Prefix Sum HashMap

Problem Statement

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

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

Intuition

If prefix[j] - prefix[i] = k, then subarray [i+1..j] sums to k. Rearranging: prefix[i] = prefix[j] - k.

Store prefix sums in HashMap as we go. For each j, count how many previous prefixes equal prefix[j]-k.

Solutions

C++

int subarraySum(vector<int>& nums, int k) {
    unordered_map<int,int> cnt; cnt[0]=1;
    int prefix=0, ans=0;
    for(int n:nums){prefix+=n;ans+=cnt[prefix-k];cnt[prefix]++;}
    return ans;
}

Java

public int subarraySum(int[] nums, int k) {
    Map<Integer,Integer> cnt=new HashMap<>(); cnt.put(0,1);
    int prefix=0, ans=0;
    for(int n:nums){prefix+=n;ans+=cnt.getOrDefault(prefix-k,0);cnt.merge(prefix,1,Integer::sum);}
    return ans;
}

JavaScript

var subarraySum = function(nums, k) {
    const cnt=new Map([[0,1]]); let prefix=0, ans=0;
    for(const n of nums){prefix+=n;ans+=(cnt.get(prefix-k)||0);cnt.set(prefix,(cnt.get(prefix)||0)+1);}
    return ans;
};

Python

from collections import defaultdict

def subarraySum(nums: list[int], k: int) -> int:
    cnt = defaultdict(int)
    cnt[0] = 1
    prefix = ans = 0
    for n in nums:
        prefix += n
        ans += cnt[prefix - k]
        cnt[prefix] += 1
    return ans

Complexity

Time: O(n), Space: O(n)