DSA Subarray Sum Equals K

The problem

Subarray Sum Equals K asks for a count. Count the subarrays whose values add up to a target k.

A subarray is a run of neighbouring values inside a list. It has a start and an end, and it cannot skip a value. The empty run does not count.

The input is a list of integers and one integer k. Values may be negative. The output is one integer, the number of matching subarrays.

cases = [
    ([1, 1, 1], 2, 2),
    ([1, 2, 3], 3, 2),
    ([1, -1, 1, -1, 1], 0, 6),
    ([3, 4, 7, 2, -3, 1, 4, 2], 7, 4),
]
for values, k, expected in cases:
    print("values", values, "k", k, "count:", expected)

Read one row at a time. In [1, 1, 1] with k = 2, the matching runs are the first two values and the last two values. In [1, 2, 3] with k = 3, the matches are [1, 2] and [3].

The third row is the hard one. Its values cancel each other out. Six different runs sum to 0, and any correct method has to find all six.

Intuition: the sliding window is a trap here

A sliding window is a window that grows at the right edge and shrinks at the left edge. It is a tempting first idea here. Grow until the sum reaches k, then shrink while the sum is too big.

That idea rests on one hidden assumption. It assumes the sum only rises as the window grows. That holds when every value is positive.

Negative values break the assumption. A later value can pull the sum back down to k. A window that was already shrunk past cannot come back.

def window_count(values, k):
    left = 0
    total = 0
    found = 0
    for right in range(len(values)):
        total += values[right]
        while total > k and left <= right:
            total -= values[left]
            left += 1
        if total == k:
            found += 1
    return found

print("window answer:", window_count([1, -1, 1, -1, 1], 0))
print("correct answer: 6")

The window method reports 3, and the correct answer is 6. It loses runs that end after a dip. The shrink step threw away starts that a later negative value would have made valid again.

The fix has to remember every earlier position. It cannot discard a start just because the current sum looks too large.

Brute force: check every subarray

The straightforward method lists every subarray. For each start, it walks to every end. It adds the values as it goes and counts a hit whenever the running sum equals k.

The code below also counts its own work. The counter rises once for each value added into a sum. That count is the printed cost of the method.

def brute_count(values, k):
    found = 0
    adds = 0
    for start in range(len(values)):
        total = 0
        for end in range(start, len(values)):
            total += values[end]
            adds += 1
            if total == k:
                found += 1
    return found, adds

for values, k in [([1, 1, 1], 2), ([1, -1, 1, -1, 1], 0)]:
    found, adds = brute_count(values, k)
    print("values", values, "k", k, "count:", found, "adds:", adds)

The answers are right, including the hard negative case. The cost is the problem. A list of length n has about n squared over two subarrays, and this method sums each one from its start.

The waste has a clear shape. The same opening values are added again and again, once for every start before them. A better method should add each value once and remember what it has seen.

The optimal solution: count matching prefix sums

A prefix sum is the sum of all values up to one position. The prefix before the list starts is 0. One new prefix appears after each value.

The key idea is one sentence. A subarray sums to k exactly when the current prefix minus an earlier prefix equals k, so the code counts earlier prefixes equal to the current prefix minus k.

A frequency map stores those earlier prefixes. A frequency map is a dict from a prefix sum to how many times that sum has appeared. It starts as {0: 1}, which stands for the empty prefix before the list.

def count_subarrays(values, k):
    seen = {0: 1}
    running = 0
    found = 0
    for value in values:
        running += value
        found += seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return found

cases = [
    ([1, 1, 1], 2),
    ([1, 2, 3], 3),
    ([1, -1, 1, -1, 1], 0),
    ([3, 4, 7, 2, -3, 1, 4, 2], 7),
]
for values, k in cases:
    print("values", values, "k", k, "count:", count_subarrays(values, k))

All four cases from the start of the article pass. The order inside the loop matters. The lookup happens before the current prefix is recorded, so a subarray always has at least one value.

Negative values need no special case. A negative value just makes the running sum go down, and the map records that lower sum like any other.

Walkthrough: prefixes on [1, -1, 1, -1, 1], k = 0

Follow the hard case one value at a time. The running sum returns to 0 or 1 after every step. Each return is a new match, because an earlier prefix with the same value closes a run that sums to 0.

values = [1, -1, 1, -1, 1]
k = 0
seen = {0: 1}
running = 0
found = 0
print("start: running 0, seen", seen, "found", found)
for value in values:
    running += value
    found += seen.get(running - k, 0)
    seen[running] = seen.get(running, 0) + 1
    print("add", value, "-> running", running, "seen", seen, "found", found)
print("answer:", found)

Read the last step slowly. The running sum is 1, and the map already holds two 1 prefixes. Each one closes a different zero-sum run ending at the last value, so the count jumps by two at once.

That jump is why the map stores counts, not just presence. A set that only said the prefix 1 had appeared would add 1 instead of 2, and the answer would come out too small.

Complexity

Time is O(n). The scan visits each value once, and each visit does one lookup and one update in the map.

Space is O(n). In the worst case every prefix sum differs, so the map grows to hold one entry per value plus the starting 0.

def work(values, k):
    seen = {0: 1}
    running = 0
    lookups = 0
    for value in values:
        running += value
        lookups += 1
        seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return lookups

def brute_adds(n):
    adds = 0
    for start in range(n):
        for end in range(start, n):
            adds += 1
    return adds

for n in [10, 20, 40]:
    print("n", n, "optimal lookups:", work([1] * n, 2), "brute adds:", brute_adds(n))

The printed counts match the claims. Doubling n doubles the lookups. The brute-force adds grow about four times over, because they follow n squared.

Edge cases and pitfalls

A k larger than any reachable count is not an error. It simply matches nothing, and the answer is 0. The same holds when the list is empty.

All-negative lists work without any change. Their prefix sums fall step by step, and equal prefixes still mark zero-sum runs when k is 0.

def count_subarrays(values, k):
    seen = {0: 1}
    running = 0
    found = 0
    for value in values:
        running += value
        found += seen.get(running - k, 0)
        seen[running] = seen.get(running, 0) + 1
    return found

print("empty list:", count_subarrays([], 5))
print("k beyond reach:", count_subarrays([1, 2], 99))
print("all negative, k = -3:", count_subarrays([-1, -2, -3], -3))
print("single value hits:", count_subarrays([4], 4))
print("single value misses:", count_subarrays([4], 5))

One pitfall hides in the starting map. Starting with an empty map instead of {0: 1} loses every subarray that starts at the first value. The empty prefix is a real earlier prefix, and it must be counted before the scan starts.

A second pitfall is the order of lookup and record. Recording the current prefix before the lookup lets a position pair with itself. That pairing describes an empty subarray, which this problem does not count.