Rain tank swings

A buried rain tank gains water when it rains and loses it when the nursery draws off. The log holds one number per hour; the question is about stretches.

The problem

A rainwater tank under a garden centre is metered every hour. The meter does not record the level; it records the change during that hour, in litres. A wet hour is positive, an hour when the pumps run is negative, and a still hour is zero.

Count the contiguous stretches whose changes add up to exactly target. A stretch is one or more consecutive entries of the log; stretches may overlap, and each distinct pair of start and end hours counts once.

Input. changes — a list of integers, the litres gained (positive) or lost (negative) in each hour. target — an integer, the net change being looked for.

Output. An integer: how many contiguous stretches have net change exactly target.

Example.

changes = [4, -1, 3, -3, 1], target = 3   ->  3

Hours 0–1 give 4 + -1 = 3, hours 0–3 give 4 + -1 + 3 + -3 = 3, and hour 2 alone gives 3. Nothing else does.

A second example, where the same stretch boundary is reused many times:

changes = [0, 0, 0], target = 0   ->  6

Three single hours, two pairs and one triple all net zero. A method that only records whether a total has been seen reports 1 here.

Constraints.

• 0 <= len(changes) <= 2 * 10^4
• -10^4 <= changes[i] <= 10^4
• -10^8 <= target <= 10^8

Hints

Approach

Brute force

Fix a start hour, extend an end hour, keep a running sum, and count the hits. That is about n(n + 1) / 2 additions — 200 million for a 20 000-hour log, most of them re-adding a prefix already computed.

The insight

A stretch summing to target is a pair of running totals differing by target, so counting stretches is counting earlier totals equal to total - target.

Let P[j] be the sum of the first j entries. The stretch i..j sums to P[j+1] - P[i], so it hits the target exactly when P[i] = P[j+1] - target. Sweeping left to right, every prefix total already in the tally is a legal start boundary, so one lookup answers the whole question for that hour. Nothing has to be sorted, and the entries may be negative — which is why a sliding window, whose sum must grow as the window grows, is unavailable.

Algorithm

1. Start a tally with {0: 1} — the boundary before hour 0.
2. Set total = 0 and found = 0.
3. For each hourly change, add it to total.
4. Add tally.get(total - target, 0) to found.
5. Increment tally[total] by one.
6. Return found.

Complexity

Time O(n) — one pass, one average-case constant-time lookup and update per hour. Space O(n) — one tally entry per distinct prefix total.

Solution

"""Rain tank swings — count target-sum stretches with a tally of prefix totals."""

from collections import defaultdict


def solve(changes, target):
    # A stretch i..j nets `target` exactly when P[i] == P[j+1] - target, so the
    # answer at hour j is how many earlier prefixes carried that value.
    seen = defaultdict(int)
    seen[0] = 1                      # the empty prefix: the boundary before hour 0
    total, found = 0, 0
    for litres in changes:
        total += litres
        found += seen[total - target]  # look up before inserting, or target 0 self-counts
        seen[total] += 1
    return found

TESTS = [
    (([4, -1, 3, -3, 1], 3), 3),
    (([0, 0, 0], 0), 6),
    (([], 5), 0),
    (([300], 300), 1),
    (([2, 2, 2, 2], 4), 3),
    (([-5, 5, -5, 5], 0), 4),
    (([7], 3), 0),
]

Pitfalls

• Leaving out the {0: 1} seed drops every stretch that starts at hour 0: the first example reports 1 instead of 3.
• Updating the tally before the lookup counts the current prefix against itself whenever target is 0. On [0, 0, 0] that reports 9 instead of 6.
• Using a set of seen totals instead of counts reports at most one stretch per hour. On [0, 0, 0] it gives 3, not 6.
• Reaching for a sliding window. With -3 in the log the sum does not rise as the window grows, so the window skips past valid stretches.

Variants

• Reserve turnstile — the same prefix map, storing the first index a total appeared at rather than a count.
• What a hash map buys you — the lesson deriving this prefix-plus-map pattern from the double loop.