Lab 03 — Prefix Sums: Subarray Sum Equals K

Goal

Master the prefix-sum + hashmap-of-complements pattern. Deliverable solves LeetCode 560 in O(N) time, O(N) space, and you can articulate why a hashmap of prefix sums (not of values) is the right abstraction, why {0: 1} is the required base case, and why this generalizes to subarray-XOR-equals-K, subarray-sum-divisible-by-K, and friends.

Background Concepts

Prefix sum identity: sum(a[l..r]) = prefix[r+1] - prefix[l]. Reformulating “find subarrays with sum K” as “find pairs of prefix sums differing by K”. Hashmap as a complement-finder. Generalization to any group operation (XOR, mod, addition over any abelian group). Review pattern Prefix Sums and Hashing.

Interview Context

This is a rite-of-passage Medium: appears at Meta, Google, Amazon, Stripe. Naive candidates write O(N²) double loops over all subarrays. Decent candidates write a prefix-sum array then double loop over endpoints — still O(N²). Strong candidates collapse to O(N) with a hashmap-of-prefix-counts. Elite candidates immediately generalize to LC 974 (subarray sums divisible by K) and LC 525 (contiguous array — recast as prefix-balance equals zero) without prompting.

Problem Statement

Given an integer array nums and an integer K, return the total number of contiguous subarrays whose sum equals K.

Constraints

• 1 ≤ N ≤ 2 × 10^4
• -1000 ≤ nums[i] ≤ 1000
• -10^7 ≤ K ≤ 10^7
• The array can contain negative numbers (this matters — sliding window does not apply).

Clarifying Questions

1. Are values negative or non-negative? (Per constraints — both. This is the crucial clarification: with non-negatives, sliding window works in O(N); with negatives, you must use prefix sums.)
2. Are zeros allowed? (Yes, and they create multiple subarrays of the same sum; the count must reflect this.)
3. Empty subarrays — count them? (No; subarrays have ≥ 1 element. But the prefix-sum technique uses an “empty prefix” of value 0, hence {0: 1} initialization.)
4. Is K always reachable? (No assumption.)
5. Can K = 0? (Yes — counts subarrays summing to 0, including those that are entirely zero.)

Examples

Initial Brute Force

Two nested loops over (l, r), sum nums[l..r], count.

def subarray_sum_brute(nums, K):
    count = 0
    for l in range(len(nums)):
        s = 0
        for r in range(l, len(nums)):
            s += nums[r]
            if s == K: count += 1
    return count

Brute Force Complexity

Time O(N²). Space O(1). At N=2×10⁴, ~4×10⁸ ops — borderline; passes in C++ but TLEs in Python.

Optimization Path

The key reformulation: a subarray nums[l..r] sums to K iff prefix[r+1] - prefix[l] == K, i.e., prefix[l] == prefix[r+1] - K.

So as we walk r from 0 to N-1, computing the running prefix sum, we ask: “How many earlier prefix sums equal prefix - K?” This is a hashmap lookup. Each step is O(1). Total O(N).

The base case is subtle and important: before processing any element, we have prefix sum 0, “seen once”. This accounts for subarrays starting at index 0 (where the missing earlier prefix is the empty prefix of value 0).

Final Expected Approach

def subarray_sum(nums, K):
    count = 0
    prefix = 0
    seen = {0: 1}                       # empty prefix
    for x in nums:
        prefix += x
        count += seen.get(prefix - K, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

Crucial ordering: lookup before insert. Otherwise the case K == 0 over-counts: every position would match itself.

Data Structures Used

• A hashmap seen: int → int mapping each prefix-sum value to the number of times it has occurred.
• A running prefix integer.
• A running count integer.

Correctness Argument

Let p_i = sum(nums[0..i-1]) (so p_0 = 0). A subarray nums[l..r] sums to K iff p_{r+1} - p_l = K iff p_l = p_{r+1} - K.

For each r, the number of valid l ∈ [0, r] is |{l : p_l == p_{r+1} - K}|. As we iterate, seen after processing index r contains exactly {p_0, p_1, ..., p_{r+1}} with multiplicities. Looking up seen[prefix - K] before inserting prefix gives |{l ∈ [0, r] : p_l == p_{r+1} - K}| — the count of valid l for the current r.

Summing over all r gives the total count of valid subarrays. The {0: 1} initialization handles the case l == 0 (where p_0 == 0 is consulted).

Complexity

Time O(N) average (hashmap ops). Space O(N) for the hashmap (worst case: all distinct prefix sums). Worst-case time degrades to O(N²) under hash collisions on adversarial input — see Phase 1 §3.

Implementation Requirements

• Initialize seen = {0: 1} before the loop. Forgetting this causes off-by-one on subarrays starting at index 0.
• Lookup before insert. This isn’t just style — for K == 0, swapping the order miscounts trivially.
• Use 64-bit accumulator if values × N could overflow 32-bit (here, 2×10^4 × 1000 = 2×10^7, safe; but build the habit).
• Don’t precompute the prefix-sum array if you don’t need to — a single running int suffices.

Tests

• Smoke: ([1,1,1], 2) → 2.
• Unit: K = 0 cases; K larger than total sum (returns 0); single element matching K (returns 1); single element not matching K (returns 0).
• Edge: nums = [0]*N, K = 0 → N*(N+1)/2. This validates that zeros are correctly counted.
• Adversarial: alternating positives and negatives summing to K many times. Construct [1,-1,1,-1,...,1,-1] with K=0; expected count is N/2 + N/2*(N/2-1)/2 + … — easier to validate against the brute force.
• Large: N = 2×10⁴, random values; assert ms-level. Compare to brute force on a 1000-element prefix.
• Random: 100 random inputs of size ≤ 200; cross-check against brute.
• Negative K: include K negative; must work because prefix sums and prefix - K use the same arithmetic.

Follow-up Questions

• “Subarray sum divisible by K?” → key by prefix % K (taking care to normalize to non-negative for languages where % can return a negative value: ((prefix % K) + K) % K). Count pairs of equal residues.
• “Subarray with XOR equal to K?” → exact same skeleton, replace + with ^. The group operation just needs an identity and an inverse.
• “Longest subarray with sum K?” → store the first occurrence of each prefix-sum value; when you see prefix - K, the length is r+1 - first[prefix - K].
• “Contiguous array (LC 525) — equal 0s and 1s?” → re-encode 0s as -1; the problem becomes “longest subarray with sum 0”.
• “If memory is tight (cannot store hashmap)?” → if values are non-negative, sliding window in O(1) extra; otherwise no known better.
• “What if the array is updated?” → Fenwick tree for prefix sums; each update O(log N), each query O(log N) (Phase 3).

Product Extension

Anomaly detection in cumulative metric streams. Imagine network ingress where a “subarray sum equals K” query asks: “how many time windows had exactly K bytes of traffic?” — useful for detecting periodicity or replay patterns. The same prefix-hashmap pattern, run online, detects these in O(1) per event with O(window-size) memory. Extends to log-aggregation systems (CloudWatch, Datadog) where dashboards scan millions of events and need O(N) algorithms.

Language/Runtime Follow-ups

• Python: defaultdict(int) is cleaner than dict.get(...). Performance is similar.
• Java: HashMap<Long, Integer> for safety. Boxing! Each getOrDefault(prefix, 0) + 1 followed by put(prefix, ...) does up to 4 boxes; pre-boxing or merge() with Integer::sum reduces it.
• Go: map[int]int; the zero value of int is 0 so seen[prefix - K] returns 0 if absent — no need for ok check on lookups (but you do need to insert seen[0] = 1 initially).
• C++: std::unordered_map<long long, int>. Beware: unordered_map<int, int>::operator[] inserts on read of a missing key, growing the map by one for every miss. Use find() for lookup-only.
• JS/TS: prefer Map<number, number>; Object keys are coerced to strings. For very large prefix sums (> 2^53), use BigInt or strings as keys.
• Adversarial hashing: crafted inputs producing many distinct prefix sums that collide can degrade to O(N²) in languages with deterministic hash. Java’s HashMap upgrades long chains to red-black trees; Python and Go randomize; C++ does not by default.

Common Bugs

1. Missing {0: 1} initialization — off-by-one for any subarray starting at index 0. Concretely: ([3], 3) returns 0 instead of 1.
2. Insert before lookup — for K == 0, every position matches itself once, returning N spurious counts.
3. Using dict[key] instead of dict.get(key, 0) in Python — KeyError on first miss.
4. Java autoboxing penalty — silent ~3× slowdown; use primitive maps (e.g., Eclipse Collections IntIntMap) or HashMap.merge.
5. C++ unordered_map::operator[] insertion-on-read — bloats the map and ruins iteration order; use find() for read-only.
6. Negative-modulo bug in the LC 974 follow-up — -3 % 5 is -3 in C++/Java/Go but 2 in Python. Normalize: ((x % K) + K) % K.
7. Recomputing prefix - K twice — minor, but (prefix - K) should be a single local for readability.

Debugging Strategy

• Trace ([1,1,1], 2) by hand. Expected seen evolution: {0:1} → {0:1,1:1} → {0:1,1:1,2:1} → {0:1,1:1,2:1,3:1}. At each step, lookup prefix - K: at step 2, prefix=2, lookup 0: count += 1. At step 3, prefix=3, lookup 1: count += 1. Total 2.
• For K = 0 issues: trace ([0,0,0], 0). seen evolves {0:1} → (lookup 0: count+=1) → {0:2} → (lookup 0: count+=2) → {0:3} → (lookup 0: count+=3). Total 6 = 3*(3+1)/2. ✓
• Diff against brute force on 100 random small inputs.

Mastery Criteria

• Recognized “subarray with sum X” + negatives as prefix-sum + hashmap in <30 seconds.
• Articulated the {0: 1} base case before coding.
• Lookup-before-insert ordering correct on first try.
• Generalized to LC 974 (mod) and LC 525 (re-encoding) without lookup.
• Identified C++ operator[] insertion-on-read trap.
• Solved on the first try with no off-by-one bugs in 5 random small tests.