LeetCode 560. Subarray Sum Equals K - From Intuitive O(n²) to Optimized O(n) solution using HashMap(Med, Java)

Anni HuangAnni Huang
3 min read

Problem Statement

Given an integer array nums and an integer k, return the total number of continuous subarrays whose sum is equal to k.

Example:

  • Input: nums = [1,1,1], k = 2
  • Output: 2 (subarrays [1,1] at indices 0-1 and 1-2)

Two Solutions

Intuitive O(n²) Solution

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

        // Check every possible subarray
        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;
    }
}

How it works:

  • Use two nested loops to generate all possible subarrays
  • For each starting position i, extend the subarray by adding elements at position j
  • Keep a running sum and increment count when sum equals k
  • Time Complexity: O(n²)
  • Space Complexity: O(1)

Optimized Using HashMap (Prefix Sum)

class Solution {
    public int subarraySum(int[] nums, int k) {
        HashMap<Integer, Integer> prefixCount = new HashMap<>();
        prefixCount.put(0, 1);  // Base case: empty prefix sum

        int prefixSum[] = new int[nums.length];
        prefixSum[0] = nums[0];
        for(int i = 1; i < nums.length; i++){
            prefixSum[i] = nums[i] + prefixSum[i-1];
        }

        int count = 0;

        for(int i = 0; i < nums.length; i++){
            // Check if there's a previous prefix sum that makes current subarray sum = k
            if(prefixCount.containsKey(prefixSum[i] - k)){
                count += prefixCount.get(prefixSum[i] - k);
            }
            // Add current prefix sum to the map
            prefixCount.put(prefixSum[i], prefixCount.getOrDefault(prefixSum[i], 0) + 1);
        }
        return count;
    }
}

Key Insight: If prefixSum[j] - prefixSum[i-1] = k, then the subarray from index i to j has sum k. Rearranging: prefixSum[i-1] = prefixSum[j] - k

How it works:

  1. Build prefix sum array: prefixSum[i] = sum of elements from index 0 to i
  2. Use HashMap to track frequency: Store how many times each prefix sum has occurred
  3. For each position i: Check if (prefixSum[i] - k) exists in the map
    • If yes, it means there are subarrays ending at position i with sum k
    • Add the frequency of (prefixSum[i] - k) to our count
  4. Update map: Add current prefix sum to the HashMap

Example walkthrough:

  • nums = [1,1,1], k = 2
  • Prefix sums: [1,2,3]
  • At i=1: prefixSum[1] = 2, looking for 2-2 = 0 (found 1 time) → count = 1
  • At i=2: prefixSum[2] = 3, looking for 3-2 = 1 (found 1 time) → count = 2

Time Complexity: O(n) Space Complexity: O(n)

The optimized solution uses the mathematical property of prefix sums to avoid checking all subarrays explicitly, reducing time complexity from O(n²) to O(n).

0
Subscribe to my newsletter

Read articles from Anni Huang directly inside your inbox. Subscribe to the newsletter, and don't miss out.

Written by

Anni Huang
Anni Huang

I’m Anni Huang, an AI researcher-in-training currently at ByteDance, specializing in LLM training operations with a coding focus. I bridge the gap between engineering execution and model performance, ensuring the quality, reliability, and timely delivery of large-scale training projects.