Subarray Sum Equals K — Leetcode 560

Original problem link: https://leetcode.com/problems/subarray-sum-equals-k/

In this problem, we are asked to find the number of subarrays in a given array that sum to a given k.

Though we can come up with a brute force solution quite easily, it is not the most efficient approach. Since we are looking at subarray sums, we can quite intuitively think of using prefix sums and calculating combinations from that. Note that the more concise solution looks like this:

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        int sum = 0;
        int combinations = 0;
        map.put(0, 1);
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
            combinations += map.getOrDefault(sum - k, 0);
            map.put(sum, map.getOrDefault(sum, 0) + 1);
        }
        return combinations;
    }
}

Intuitively, if there exists a prefix sum of current sum — k, then there is a new combination (subarray). We also need to put in 0->1 as an initial mapping to count the combinations correctly.

Previously, I came up with my own solution without the 0->1 mapping, which is a method I found unintuitive. But without that mapping, I had to write extra cases to take care of sum = k and k = 0 cases which were not being counted correctly.

class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> map = new HashMap<>();
        int sum = 0;
        int combinations = 0;
        for (int i = 0; i < nums.length; i++) {
            sum += nums[i];
            map.put(sum, map.getOrDefault(sum, 0) + 1);

            if (sum == k) {
                // sum - k = 0, so since we didn't map 0->1 we have to take care of this case
                combinations += map.get(sum);
                if (map.containsKey(0) && k != 0) {
                    combinations += map.get(0);
                    map.put(0, map.get(0) - 1);
                }
            } else if (map.containsKey(sum - k) && k != 0) {
                combinations += map.get(sum - k);
            } else if (k == 0) {
                // otherwise, sum - 0 counts an empty array
                combinations += map.get(sum) - 1;
            }
        }
        return combinations;
    }
}

The speed is roughly the same for the two solutions — the mapping with 0->1 is sometimes a little faster.