Optimize the code · Medium

Speed Up Subarray Sum Equals K

Medium Make it fast Most askedSeen at Meta, Google, Amazon

subarraySum(nums, k) returns how many contiguous subarrays of nums add up to exactly k. Values can be negative.

The version below tries every start and end position. It is correct, but a 70,000-element input takes too long. Make it O(n) without changing the results.

Example 1

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

Example 2

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

The slow code

Every case must finish within 300 ms (budgets are per language in the editor).

function subarraySum(nums, k) {
  let count = 0;
  for (let i = 0; i < nums.length; i++) {
    let sum = 0;
    for (let j = i; j < nums.length; j++) {
      sum += nums[j];
      if (sum === k) count++;
    }
  }
  return count;
}

function subarraySum(nums: number[], k: number): number {
  let count = 0;
  for (let i = 0; i < nums.length; i++) {
    let sum = 0;
    for (let j = i; j < nums.length; j++) {
      sum += nums[j];
      if (sum === k) count++;
    }
  }
  return count;
}

def subarray_sum(nums: list[int], k: int) -> int:
    count = 0
    for i in range(len(nums)):
        total = 0
        for j in range(i, len(nums)):
            total += nums[j]
            if total == k:
                count += 1
    return count

<?php

function subarraySum(array $nums, int $k): int
{
    $count = 0;
    $n = count($nums);
    for ($i = 0; $i < $n; $i++) {
        $sum = 0;
        for ($j = $i; $j < $n; $j++) {
            $sum += $nums[$j];
            if ($sum === $k) {
                $count++;
            }
        }
    }
    return $count;
}

class Solution {
    public int subarraySum(int[] nums, int k) {
        int count = 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) count++;
            }
        }
        return count;
    }
}

Hints

Hint 1

If prefix[j] - prefix[i] = k, then the subarray i..j-1 sums to k.

Hint 2

Keep a hash map counting how often each prefix sum has occurred so far.

Make it faster

Loading the editor… (JavaScript is required to run code in your browser.)

Solution and explanation

Show the solution

Time O(n) · Space O(n)

Why it is slow: two nested loops visit all n(n+1)/2 subarrays: about 2.5 billion for 70,000 elements.

The fix: walk once with a running sum s. A subarray ending here sums to k exactly when some earlier prefix equals s - k. Keep a hash map of how many times each prefix sum has appeared (starting with {0: 1} for the empty prefix), add count[s - k] to the answer, then record s. O(n) time.

Sliding windows do not work here because values can be negative; saying why is a good signal.

Optimized solution

function subarraySum(nums, k) {
  const seen = new Map([[0, 1]]); // prefix sum -> how many times seen
  let sum = 0;
  let count = 0;
  for (const x of nums) {
    sum += x;
    count += seen.get(sum - k) ?? 0;
    seen.set(sum, (seen.get(sum) ?? 0) + 1);
  }
  return count;
}

function subarraySum(nums: number[], k: number): number {
  const seen = new Map<number, number>([[0, 1]]); // prefix sum -> how many times seen
  let sum = 0;
  let count = 0;
  for (const x of nums) {
    sum += x;
    count += seen.get(sum - k) ?? 0;
    seen.set(sum, (seen.get(sum) ?? 0) + 1);
  }
  return count;
}

from collections import Counter

def subarray_sum(nums: list[int], k: int) -> int:
    seen = Counter({0: 1})  # prefix sum -> how many times seen
    total = count = 0
    for x in nums:
        total += x
        count += seen[total - k]
        seen[total] += 1
    return count

<?php

function subarraySum(array $nums, int $k): int
{
    $seen = [0 => 1]; // prefix sum => how many times seen
    $sum = 0;
    $count = 0;
    foreach ($nums as $x) {
        $sum += $x;
        $count += $seen[$sum - $k] ?? 0;
        $seen[$sum] = ($seen[$sum] ?? 0) + 1;
    }
    return $count;
}

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

Common mistakes

Forgetting the initial {0: 1}: subarrays starting at index 0 are missed.
Recording the current prefix before looking up s - k (wrong when k = 0).
Using a sliding window: only valid when all values are positive.

Follow-up questions

Return the longest such subarray instead of the count.
Count subarrays whose sum is divisible by k.

Speed Up Contains Duplicate · Easy
Speed Up Two Sum · Easy
Speed Up Recursive Fibonacci · Easy
Speed Up the Most Frequent Word · Easy
Speed Up the Active-Customer Join · Easy
Speed Up Range Sum Queries · Medium