560. Subarray Sum Equals K

Problem Statement

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

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

Example 1:

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

Constraints:

• 1 <= nums.length <= 2 * 10<sup>4</sup>

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

• -10<sup>7</sup> <= k <= 10<sup>7</sup>

Solution

This is a follow up question to the Longest Subarray With Sum K (link)

Brute Force Approach

Create all subarrays using a nested loop structure, calculate their sums, and increment the count if the sum matches the target sum.

Time - O(n^2)

Space - O(1)

class Solution {
    public int subarraySum(int[] nums, int k) {
        int ans = 0;

        for(int i=0; i<nums.length; i++){
            int sum =0;
            for(int j=i; j<nums.length; j++){
                sum+=nums[j];
                if(sum==k) ans+=1;
            }
        }
        return ans;
    }
}

Optimal Approach

We employ a prefix sum approach to compute the sum of any subarray(link). The sums are stored in a map along with their respective frequencies. Our strategy involves consistently determining the difference between the target and the prefix sum. We then verify whether this difference has been encountered in the prefix sum map. If so, we increment the answer count by the frequency associated with the difference in the prefix sum.

Note: It is crucial to include the prefix sum of 0 with a count of 1 in the map. This inclusion is necessary because if the target matches the current prefix sum exactly, the difference will be 0. Although 0 might not be present in the map, it signifies a valid subarray.

Time - O(n)

Space - O(1)

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

        Map<Integer, Integer> prefixSum = new HashMap<>();

        prefixSum.put(0, 1);

        int sum = 0;

        int ans = 0;

        for(int i=0; i<nums.length; i++){
            sum+=nums[i];
            int diff = sum-k;

            if(prefixSum.containsKey(diff)){
                ans+=prefixSum.get(diff);
            }

            if(!prefixSum.containsKey(sum)){
                prefixSum.put(sum,0);
            }

            prefixSum.put(sum, prefixSum.get(sum)+1);
        }

        return ans;
    }
}