Count target-sum subarrays: differences of prefixes

Turn a range question into a historical lookup

A correct baseline starts a running sum at each left endpoint and extends right, counting matches. This costs O(n²) time and O(1) auxiliary space. A sliding window is not generally valid: removing a negative value increases the sum, and adding a negative can repair a sum that was too large. Instead define prefix P[j] as the sum of the first j values, with P[0] = 0. Subtracting cancels the shared prefix: the sum over [i, j) is P[j] - P[i].

Rearrange the desired equality to P[i] = P[j] - target. A frequency map, not a set, records how many earlier endpoints qualify. Before processing prefix j, the map counts exactly positions 0 through j-1, and the answer counts all matches ending before j. Add the matching frequency, then insert the new prefix. Inserting first would count an empty range whenever the target is zero. The seeded zero represents a legitimate range beginning at input position zero; it is not an arbitrary special case.

Each range has one unique right endpoint, so this procedure counts every valid range exactly once. Expected time is O(n), auxiliary space O(n) for distinct prefix sums, and output O(1) integers. Python integers avoid fixed-width overflow; these bounds count arithmetic operations rather than bit-level cost on enormous integers. The answer itself can be n(n+1)/2, as an all-zero input demonstrates.

Follow-up 1: count only ranges of length at most two

Predict which prefix frequencies must expire. At endpoint j, allowable starts are max(0, j-2) through j-1; a permanent map overcounts long ranges.

Keep a queue of eligible prefix sums, decrementing frequencies as old positions expire. Delete zero counts. This yields expected O(n) time and O(min(n, limit)) state. Expiring a value entirely when only one of several copies leaves is wrong.

Follow-up 2: find a longest matching range

Now frequency is the wrong return information: only the earliest equal prefix maximizes length for the current endpoint. A senior candidate derives the equation and changes map values to match the requested output. A lead candidate defines integer limits, retention, and whether late ledger corrections require recomputing historical answers; this append-only method does not support edits.

Run and check

From the repository root:

cd curriculum/01-code/02-data-structures-algorithms/problems/06-subarray-sum-count
python -m unittest -v test_solution.py

Reference implementation (download file, source below) · Contract and oracle tests (download file, source below). Read the tests after your attempt. A green reference suite verifies the supplied implementation; it does not demonstrate independent transfer. Reimplement one follow-up with the reference closed and explain which old invariant no longer holds.

def _integers(values):
    if not isinstance(values, (list, tuple)) or any(type(x) is not int for x in values):
        raise ValueError("expected a list or tuple of integers, excluding bool")

def subarray_sum_count(nums, target):
    _integers(nums)
    if type(target) is not int:
        raise ValueError("target must be an integer")
    prefix_counts = {0: 1}
    prefix = total = 0
    for value in nums:
        prefix += value
        total += prefix_counts.get(prefix - target, 0)
        prefix_counts[prefix] = prefix_counts.get(prefix, 0) + 1
    return total

import itertools
import unittest
from solution import subarray_sum_count

class PrefixTests(unittest.TestCase):
    def test_boundaries(self):
        self.assertEqual(subarray_sum_count([1, -1, 0], 0), 3)
        self.assertEqual(subarray_sum_count([], 0), 0)
        self.assertEqual(subarray_sum_count([0] * 100, 0), 5050)
        self.assertEqual(subarray_sum_count([10**100, -(10**100)], 0), 1)
        for nums, target in [([True], 1), ([1.5], 1), ([1], False), (None, 0)]:
            with self.assertRaises(ValueError):
                subarray_sum_count(nums, target)

    def test_range_sum_oracle(self):
        for n in range(6):
            for nums in itertools.product(range(-1, 2), repeat=n):
                for target in range(-2, 3):
                    expected = sum(sum(nums[i:j]) == target
                                   for i in range(n) for j in range(i + 1, n + 1))
                    self.assertEqual(subarray_sum_count(nums, target), expected)

if __name__ == "__main__":
    unittest.main()