Subarray Sum Equals K

NeetCode: Subarray Sum Equals K · LeetCode: 560. Subarray Sum Equals K · Difficulty: Medium

Given an integer array nums and an integer k, return the number of contiguous non-empty subarrays whose sum equals k. Elements may be negative, so a sliding window, which needs non-negative values, does not apply.

Example

Text

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

Brute force: every start, extend the end

For each start index, keep a running long sum while extending the end, and count each time the sum equals k. That checks all n(n+1)/2 subarrays.

Java

public class Main {
    static class Solution {
        public int subarraySum(int[] nums, int k) {
            int count = 0;
            for (int i = 0; i < nums.length; i++) {
                long sum = 0;
                for (int j = i; j < nums.length; j++) {
                    sum += nums[j];
                    if (sum == k) {
                        count++;
                    }
                }
            }
            return count;
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        System.out.println(sol.subarraySum(new int[] {1, 1, 1}, 2));
        System.out.println(sol.subarraySum(new int[] {1, 2, 3}, 3));
        System.out.println(sol.subarraySum(new int[] {1, -1, 0}, 0));
        System.out.println(sol.subarraySum(new int[] {}, 0));
    }
}

Output

Time: O(n^2). Space: O(1).

Optimal: prefix sums and a hash map

Let prefix[j] be the sum of the first j elements. The subarray nums[i..j] sums to k exactly when prefix[j] - prefix[i] == k, that is prefix[i] == prefix[j] - k. So while scanning, keep a map from each prefix sum seen so far to how many times it occurred. At each position, add map[current - k] to the answer, then record the current prefix. Seed the map with {0: 1} so subarrays that start at index 0 are counted (their "earlier prefix" is the empty prefix, 0). Negatives are handled because nothing assumes the prefix only grows.

Walkthrough for [1, 2, 3], k = 3:

x	prefix	need (prefix - k)	found	map afterwards
start	0			{0: 1}
1	1	-2	0	{0: 1, 1: 1}
2	3	0	1	{0: 1, 1: 1, 3: 1}
3	6	3	1	{0: 1, 1: 1, 3: 1, 6: 1}

Total is 2 ([1, 2] and [3]).

Java

import java.util.HashMap;
import java.util.Map;

public class Main {
    static class Solution {
        public int subarraySum(int[] nums, int k) {
            Map<Long, Integer> seen = new HashMap<>(); // prefix sum -> times seen
            seen.put(0L, 1);
            long prefix = 0;
            int count = 0;
            for (int x : nums) {
                prefix += x;
                count += seen.getOrDefault(prefix - k, 0);
                seen.merge(prefix, 1, Integer::sum);
            }
            return count;
        }
    }

    public static void main(String[] args) {
        Solution sol = new Solution();
        System.out.println(sol.subarraySum(new int[] {1, 1, 1}, 2));
        System.out.println(sol.subarraySum(new int[] {1, 2, 3}, 3));
        System.out.println(sol.subarraySum(new int[] {1, -1, 0}, 0));
        System.out.println(sol.subarraySum(new int[] {}, 0));
    }
}

Output

Time: O(n) expected. Space: O(n).

The order matters: look up first, then insert the current prefix, otherwise k = 0 would wrongly count an empty subarray. The prefix sum is a long because an int running sum can wrap around on large inputs. The map is keyed by Long, and 0L and prefix - k are long values, so lookups use Long.equals and never ==. getOrDefault returns 0 for a missing prefix, and merge(prefix, 1, Integer::sum) adds one to the stored count.