☕ DSA Java Articles

Topics › Arrays › Medium

Count subarrays with given sum

Medium TUF ID 561 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Problem

Given an integer array nums and an integer k, return the number of contiguous subarrays whose elements sum to exactly k. A subarray is a non-empty, contiguous slice nums[i..j] with i ≤ j. Elements may be negative or zero, so the same value can be reached by many different slices, and a slice can “overshoot” k and come back to it.

Constraints: 1 ≤ n ≤ 2 × 10^4, -1000 ≤ nums[i] ≤ 1000, -10^7 ≤ k ≤ 10^7 on LeetCode. The code handles n = 0 and uses long for sums so that larger values are also safe.

Example 1

nums = [1, 1, 1], k = 2
answer = 2               // [1, 1] at indices 0..1 and at indices 1..2

Example 2

nums = [1, 2, 3], k = 3
answer = 2               // [1, 2] and [3]

Example 3

nums = [1, -1, 0], k = 0
answer = 3               // [1, -1], [0] and [1, -1, 0]

Because negatives are allowed, the sliding-window technique that works for positive arrays does not apply here: growing the window can make the sum go down, so there is no monotonic quantity to slide on.

Approach 1: Brute force

Intuition. A subarray is determined by its two endpoints. Enumerate every pair (i, j) with i ≤ j, add up the elements in between, and count the pairs where the sum is k. This is the definition written as code and is the baseline the faster methods must agree with.

Algorithm.

  1. For every start i and every end j ≥ i:
  2. Compute sum = nums[i] + ... + nums[j] with a third loop.
  3. If sum == k, increment the counter.
static int bruteForce(int[] nums, int k) {
    int n = nums.length, count = 0;
    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            long sum = 0;
            for (int p = i; p <= j; p++) sum += nums[p];
            if (sum == k) count++;
        }
    }
    return count;
}

Complexity. Time O(n^3): O(n^2) pairs, each summed in O(n). Space O(1).

Approach 2: Better, running sum per start index

Intuition. The inner summation repeats work: the sum of nums[i..j] is the sum of nums[i..j-1] plus nums[j]. So for a fixed start i, extend the end j one step at a time and maintain the sum incrementally. This removes one whole loop.

Algorithm.

  1. For every start i: set sum = 0.
  2. For every end j from i to n - 1: sum += nums[j]; if sum == k, increment the counter.
static int better(int[] nums, int k) {
    int n = nums.length, count = 0;
    for (int i = 0; i < n; i++) {
        long sum = 0;
        for (int j = i; j < n; j++) {
            sum += nums[j];
            if (sum == k) count++;
        }
    }
    return count;
}

Complexity. Time O(n^2): each pair (i, j) is now processed in constant time. Space O(1).

Note that we cannot break when sum > k, because a later negative element could bring the sum back down to k. That is precisely why this problem is harder than its all-positive cousin.

Approach 3: Optimal, prefix sums and a hash map

Intuition. Let prefix[j] be the sum of the first j elements (prefix[0] = 0). The sum of the subarray nums[i..j-1] is prefix[j] - prefix[i]. So the question “how many subarrays ending just before j sum to k?” becomes “how many earlier prefixes prefix[i] equal prefix[j] - k?” That is a counting question about values we have already seen, and a hash map from prefix value to its number of occurrences answers it in O(1).

Walk left to right, maintaining the running prefix. At each step, look up how many times prefix - k has appeared, add that to the answer, then record the current prefix in the map. The initial entry freq[0] = 1 represents the empty prefix, which is what lets subarrays that start at index 0 be counted.

This is the same “have I seen the complement before?” idea as in Two Sum, applied to prefix sums instead of raw values. Because the map counts all earlier occurrences, repeated prefix values (which happen whenever a stretch sums to zero) are handled correctly.

Algorithm.

  1. Create freq with freq[0] = 1; set prefix = 0, count = 0.
  2. For each x: prefix += x; count += freq[prefix - k] (zero if absent); freq[prefix]++.
  3. Return count.
static int optimal(int[] nums, int k) {
    Map<Long, Integer> freq = new HashMap<>();   // prefix sum -> how many times it has occurred so far
    freq.put(0L, 1);                             // the empty prefix, so subarrays starting at index 0 are counted
    long prefix = 0;
    int count = 0;
    for (int x : nums) {
        prefix += x;
        count += freq.getOrDefault(prefix - k, 0);
        freq.merge(prefix, 1, Integer::sum);
    }
    return count;
}

Dry run on nums = [1, -1, 0], k = 0:

x prefix prefix - k freq[prefix - k] count freq after
start 0     0 {0: 1}
1 1 1 0 0 {0: 1, 1: 1}
-1 0 0 1 1 {0: 2, 1: 1}
0 0 0 2 3 {0: 3, 1: 1}

The three counted subarrays are [1, -1] (prefix 0 at step 2 matched the empty prefix), and [0] plus [1, -1, 0] (prefix 0 at step 3 matched both earlier zeros).

Complexity. Time O(n): one pass with O(1) average-time map operations. Space O(n) for the map in the worst case (all prefixes distinct).

Why must we look up before inserting the current prefix? If we inserted first, then for k = 0 the current prefix would match itself and count an empty subarray, which is not allowed. Insert after the lookup, exactly as Two Sum inserts after checking.

Edge cases and pitfalls