DP with hash maps

Many problems ask whether some contiguous piece of a sequence hits a target sum or a target average. When the natural state is "have I seen this exact prefix sum before, and how many times", a dictionary keeps that state compact, even when the possible sums are far too large to index with a plain array.

Take a sequence of numbers and ask how many contiguous pieces of it have a given average. A direct check of every possible piece costs $O(N^2)$. Tracking prefix sums in a hash map answers the same question in $O(N)$.

The key insight: prefix sum as state

A piece of the sequence from position l to position r sums to prefix[r] - prefix[l - 1], where prefix[i] is the sum of the first i elements. Two positions with the same prefix sum mean the piece between them sums to exactly zero. That single fact turns "does some subarray sum to zero" into "has this prefix sum value appeared before", which a hash map answers in constant time.

If the sums involved can be as large as $10^9$, an array indexed directly by sum is not an option. A hash map only ever holds as many entries as there are distinct prefix sums actually seen, which is at most $N + 1$ for a sequence of length $N$.

Mean-shift: turning "average" into "sum equals zero"

A subarray's average equals a target value exactly when the sum of that subarray, after subtracting the target from every one of its elements, comes out to zero. Subtracting a constant target from every element of the sequence shifts a subarray's sum by target times the subarray's length; a subarray whose original average was exactly target has a shifted sum of precisely 0, since its original sum was already target * length.

This turns "find every subarray whose average is target" into "find every subarray of the shifted sequence whose sum is zero", exactly the zero-sum question the prefix-sum hash map already answers.

Counting subarrays with a hash map

Keep a dictionary mapping each prefix sum value seen so far to how many times it has occurred, starting with the empty prefix, sum 0, occurring once. Walk through the shifted sequence, updating a running prefix sum at each step. Before recording the current prefix sum, look up how many earlier positions shared this exact value; each one marks a subarray, ending here, that sums to zero. Then increment the count for this prefix sum, so later positions can find it too.

1import sys2
3
4def main() -> None:5    input_data = sys.stdin.read().split()6    n = int(input_data[0])7    a = [int(input_data[i + 1]) for i in range(n)]8    target = int(input_data[n + 1])9
10    # Mean-shift: a subarray's average equals target exactly when the sum of11    # (element - target) over that subarray is 0. Counting zero-sum12    # subarrays only needs a hash map keyed by prefix sum, since two equal13    # prefix sums mean the subarray between them sums to zero.14    shifted = [x - target for x in a]15
16    count_by_prefix = {0: 1}  # the empty prefix, before any elements17    prefix = 018    answer = 019
20    for val in shifted:21        prefix += val22        answer += count_by_prefix.get(prefix, 0)23        count_by_prefix[prefix] = count_by_prefix.get(prefix, 0) + 124
25    print(answer)26
27
28if __name__ == "__main__":29    main()

6
4 7 1 3 5 2
3

2

With a = [4, 7, 1, 3, 5, 2] and a target average of 3, subtracting 3 from every element gives [1, 4, -2, 0, 2, -1]. The running prefix sums, including the empty prefix before any element, are 0, 1, 5, 3, 3, 5, 4. The value 3 appears twice, at the prefix after 3 elements and again after 4 elements, so the single element between them, index 3 with value 3, is a subarray of length 1 averaging exactly 3. The value 5 also appears twice, after 2 elements and again after 5 elements, so the three elements between them, indices 2 through 4 with values 1, 3, 5, form a subarray summing to 9 over 3 elements, an average of exactly 3. Those are the only two repeated prefix sums, so the program reports 2 matching subarrays.

Why this runs in O(N) time

A standard DP over every possible sum would need an array sized to the largest sum the problem allows, which is not feasible once that range reaches into the billions. A hash map instead grows only with the number of distinct prefix sums actually produced while scanning the sequence once, at most $N + 1$ of them. Each lookup and update is $O(1)$ on average, so the whole scan costs $O(N)$.

Handling repeated prefix sums

The idiom for updating a count safely, whether or not the key has been seen yet, is count_by_prefix[prefix] = count_by_prefix.get(prefix, 0) + 1, which works whether this is the first time this prefix sum has appeared or the tenth. When the goal is a best cost rather than a count, best[key] = min(best.get(key, float("inf")), new_cost) follows the same shape, updating only when the new value actually improves on whatever was stored.

Common mistakes

Forgetting the empty prefix is the most common error. Position -1, before any elements at all, has a prefix sum of 0 and must be counted once from the start; skipping it misses every subarray that begins at index 0.

Getting the mean-shift arithmetic backward is another. Subtracting the target from every element, not the other way around, is what makes a subarray's shifted sum equal zero exactly when its original average equals the target; double-check the sign on a small example by hand before trusting it on the full input.

A third mistake is reaching for floating-point sums when every value involved is an integer. The target average and every array element here are integers, so the shifted values and every prefix sum stay exact integers too; there is no reason to introduce floating-point rounding into a computation that integers already handle exactly.

Recap

A DP whose state is "which prefix sum have I seen" belongs in a hash map, not an array, whenever the range of possible sums is too large to index directly. Mean-shift converts a target-average question into the zero-sum question a prefix-sum hash map already answers, and the same counting idiom extends to tracking a best cost per prefix sum instead of just a count.