Data Structures & Algorithms (Prefix Sum): Subarray Sum Equals K

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

Example 1:

Input: nums = [1, -1, 1, -1, 2, 3, 3, 4, 1, 5, 2], k = 8
Output: 5
Explanation: [1, -1, 1, -1, 2, 3, 3], [1,-1, 2, 3, 3], [2, 3, 3], [3, 4, 1] and [1, 5, 2] are all possible subarrays whose sum = 8.

Playground: Leetcode 560

General Observations

Q. Can nums contain negative integers?
A. Yes, nums can contain negative integers. For example, nums = [-1, 0, 1, 2] is a valid input.

Approach 1: Nested Loops (Brute Force)

Use nested loops to iterate over all possible subarrays, maintain a running sum for each, and increment the count whenever the sum = k.

Algorithm

public int subarraySum(int[] nums, int k) {
    
    int count = 0;

    for(int i = 0; i < nums.length; i++) {
        int currentSum = 0;
        for(int j = i; j < nums.length; j++) {
            currentSum += nums[j];
            if(currentSum == k) {
                count++;
            }
        }
    }

    return count;

}

Visualization

Time Complexity

The operations inside the inner loop take constant time. Therefore, the overall time complexity depends on how many times the inner loop executes.

For every i, the inner loop runs (n-i) times. Since, i ranges from [0, n-1], total executions = n + (n-1) + (n-2) + ... + 1 = n(n+1)/2.

Since, the total execution's count grow by the order of n², the overall time complexity = O(n²).

Space Complexity

Since we are not using any additional data structures and only a few variables (count, currentSum), the extra space used does not grow with input size.

Hence, the overall space complexity = O(1).

Key observations

The sum of any subarray can be expressed as the difference of two prefix sums. For example, in nums = [1, -1, 1, -1, 2, 3, 3, 4, 1, 5, 2], the sum of the subarray [3, 4, 1] can be obtained by (prefixSum[8] - prefixSum[5]).

Also, for any index i, there can be multiple subarrays ending at i whose sum equals k. For example, in nums = [1, -1, 1, -1, 2, 3, 3, 4, 1, 5, 2], the subarrays [1, -1, 1, -1, 2, 3, 3], [1,-1, 2, 3, 3], and [2, 3, 3] all end at index 6 and have a sum of 8.

Approach 2: Prefix Sum

Iterate through nums while maintaining a running currentPrefixSum and a map to store the frequency of prefix sums encountered so far.

At any index = i, if a prefix sum equal to (currentPrefixSum - k) has been seen before, it means there are frequency[currentPrefixSum - k] subarrays ending at index i whose sum equals k.

Algorithm

public int subarraySum(int[] nums, int k) {

    Map<Integer, Integer> frequencyPrefixSumMap = new HashMap<>();
    frequencyPrefixSumMap.put(0, 1); // handles subarrays starting at index 0 whose sum equals k.

    int count = 0;
    int currentPrefixSum = 0;
    for(int num: nums) {
        currentPrefixSum += num;

        if(frequencyPrefixSumMap.containsKey(currentPrefixSum - k)) {
            count += frequencyPrefixSumMap.get(currentPrefixSum - k);
        }

        frequencyPrefixSumMap.put(currentPrefixSum, frequencyPrefixSumMap.getOrDefault(currentPrefixSum, 0) + 1);
    }

    return count;
}

Visualization

Time Complexity

We traverse the array exactly once. Since for each element, we perform constant-time operations (update currentPrefixSum, insert / update map), the overall time complexity = O(n).

Space Complexity

In the worst-case scenario, all prefix sums are unique. Since, we store each prefix sum in the map, the map can grow upto size n. Hence, the overall space complexity = O(n).