{"id":"dsa-complexity-prefix-sums","slug":"prefix-sums","title":"Prefix sums","text":"A step counter keeps how many steps you walked each day. How many steps from day 3 to day 10? `sum(steps[3:11])` answers in O(k) for a range of k days. That's fine once. But an app that answers 100,000 such questions about 100,000 days adds up billions of numbers, the same ones again and again.\n\n### Prefix sums\n\nDo the adding once. A **prefix sum** list holds, at index `i`, the sum of the first `i` values:\n\n```text\nvalues:     5   2   7   1   4\nprefix:  0  5   7  14  15  19\n```\n\n`prefix[0]` is 0, the sum of no values, so the list is one longer than `values`. The sum of `values[start:end]` is then `prefix[end] - prefix[start]`: `prefix[end]` adds up everything before `end`, `prefix[start]` everything before `start`, and taking one from the other leaves exactly the values from `start` to `end - 1`. One subtraction, O(1), for any range:\n\n```python\ndef make_prefix(values):\n    prefix = [0]\n    for value in values:\n        prefix.append(prefix[-1] + value)\n    return prefix\n\n\nvalues = [5, 2, 7, 1, 4]\nprefix = make_prefix(values)\nprint(prefix)\nprint(prefix[4] - prefix[1], sum(values[1:4]))\nprint(prefix[5] - prefix[0], sum(values))\n```\n\n```output\n[0, 5, 7, 14, 15, 19]\n10 10\n19 19\n```\n\nBuilding it is O(n) time and O(n) memory, once. After that, every range sum is O(1). The `start` and `end` work like a slice's, so they're easy to remember. The common mistake is to leave out the 0 at the front: then `prefix[i]` includes `values[i]`, every index is off by one, and a range that starts at 0 needs a special case.\n\n### `itertools.accumulate`\n\nThe standard library makes the running totals for you, in C:\n\n```pycon\n>>> from itertools import accumulate\n>>> list(accumulate([5, 2, 7, 1, 4]))\n[5, 7, 14, 15, 19]\n>>> list(accumulate([5, 2, 7, 1, 4], initial=0))\n[0, 5, 7, 14, 15, 19]\n```\n\nWithout `initial=0` you get the totals without the 0 at the front; with it, the prefix list above.\n\n### How many stretches add up to k?\n\nA harder question: how many stretches of values in a row (**subarrays**) add up to exactly `k`? Trying every start and end is O(n²). Prefix sums turn it into lookups. The stretch from `start` to `end` adds up to `k` when `prefix[end] - prefix[start] == k`, that is, when `prefix[start] == prefix[end] - k`. So walk along the list keeping the running total, and at each value count how many earlier prefix sums were `total - k`. A dict of counts answers that in O(1):\n\n```python\ndef count_stretches(values, k):\n    seen = {0: 1}\n    total = 0\n    count = 0\n    for value in values:\n        total += value\n        if total - k in seen:\n            count += seen[total - k]\n        if total not in seen:\n            seen[total] = 0\n        seen[total] += 1\n    return count\n\n\nprint(count_stretches([3, 4, -2, 5, 1, -4, 6], 7))\n```\n\n```output\n3\n```\n\nThe three stretches are 3, 4, then 4, -2, 5, and then 3, 4, -2, 5, 1, -4. `seen` maps each prefix sum to how many times it has come up so far, and it starts as `{0: 1}`: that's the empty prefix before the first value, so stretches that start at index 0 are counted too. Each value costs O(1) on average, so the function is O(n) time and O(n) memory. It works with negative values as well, which the sliding windows of the next step can't handle.\n\n### The biggest sum of a stretch\n\nAnother classic: which stretch of values in a row has the biggest sum, when values can be negative? The sum of a stretch is `prefix[end] - prefix[start]`, so this is much like `best_profit` from the first step, run on the prefix sums: remember the smallest prefix sum so far. Only much like, though: `best_profit` never gives less than 0, as if you could always buy and sell on the same day. On prefix sums that's a stretch of no values at all, so it gives 0 when every value is negative. The biggest sum is usually found another way, known as **Kadane's algorithm**: keep the best sum of a stretch that ends at the current value. That stretch either starts fresh at the value or adds the value to the best stretch that ended just before it, whichever is bigger. For a list that isn't empty:\n\n```python\ndef biggest_stretch(values):\n    best = values[0]\n    ending = 0\n    for value in values:\n        ending = max(value, ending + value)\n        best = max(best, ending)\n    return best\n\n\nprint(biggest_stretch([3, -5, 4, -1, 2, 1, -6, 2]))\nprint(biggest_stretch([-4, -2, -7]))\n```\n\n```output\n6\n-2\n```\n\nThe best stretch is 4, -1, 2, 1. At the 4, the best stretch ending just before it adds up to -2, so starting fresh is better, and the 3 and the -5 are left behind. When every value is negative, the best stretch is the biggest value on its own, -2, which is why `best` starts at the first value and not at 0. One pass, O(n) time and O(1) memory.\n","task":"`minutes` holds how many minutes you practiced the piano each day, and `practiced(first_day, last_day)` should give the minutes from `first_day` to `last_day`, both days included. Fill in the blank with the right index of `prefix`. The program should print:\n\n```\n75\n140\n20\n```\n","kind":"fill","seed":"from itertools import accumulate\n\nminutes = [20, 45, 0, 30, 15, 60, 10, 25]\nprefix = list(accumulate(minutes, initial=0))\n\n\ndef practiced(first_day, last_day):\n    return prefix[[[last_day + 1‖1 + last_day]]] - prefix[first_day]\n\n\nprint(practiced(1, 3))\nprint(practiced(3, 7))\nprint(practiced(0, 0))\n","solution":"from itertools import accumulate\n\nminutes = [20, 45, 0, 30, 15, 60, 10, 25]\nprefix = list(accumulate(minutes, initial=0))\n\n\ndef practiced(first_day, last_day):\n    return prefix[last_day + 1] - prefix[first_day]\n\n\nprint(practiced(1, 3))\nprint(practiced(3, 7))\nprint(practiced(0, 0))\n","hints":["The days from `first_day` to `last_day`, both included, are the slice `minutes[first_day:last_day + 1]`.","The sum of `values[start:end]` is `prefix[end] - prefix[start]`, and here `end` is one past the last day: `last_day + 1`."],"tests":[{"name":"Output","stdin":"","expect":"75\n140\n20","hidden":false}],"require":[],"forbid":[],"more":[{"task":"Write `range_sums(values, queries)`. Each query is a tuple `(start, end)`, and its answer is the sum of `values[start:end]`. Return a list of the answers, in the order of the queries. The hidden test has 200,000 values and 100,000 queries: build the prefix sums once and answer each query in O(1).\n\nFor example, `range_sums([5, 2, 7, 1, 4], [(1, 4), (0, 5), (2, 3), (3, 3)])` returns:\n\n```\n[10, 19, 7, 0]\n```\n","kind":"code","seed":"def range_sums(values, queries):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(range_sums([5, 2, 7, 1, 4], [(1, 4), (0, 5), (2, 3), (3, 3)]))\n","solution":"from itertools import accumulate\n\n\ndef range_sums(values, queries):\n    prefix = list(accumulate(values, initial=0))\n    return [prefix[end] - prefix[start] for start, end in queries]\n\n\nif __name__ == \"__main__\":\n    print(range_sums([5, 2, 7, 1, 4], [(1, 4), (0, 5), (2, 3), (3, 3)]))\n","hints":["First make the prefix sums, once, before you look at the queries: with a loop, or with `accumulate(values, initial=0)`.","For each query `(start, end)`, the answer is `prefix[end] - prefix[start]`. A loop over the queries can unpack each tuple: `for start, end in queries:`."],"tests":[{"name":"range_sums([5, 2, 7, 1, 4], [(1, 4), (0, 5), (2, 3), (3, 3)])","stdin":"","call":"range_sums([5, 2, 7, 1, 4], [(1, 4), (0, 5), (2, 3), (3, 3)])\n","expect":"[10, 19, 7, 0]","hidden":false,"returns":true},{"name":"range_sums([], [(0, 0)])","stdin":"","call":"range_sums([], [(0, 0)])\n","expect":"[0]","hidden":true,"returns":true},{"name":"range_sums([-3, 8, -1], [(0, 2), (1, 3), (0, 3)])","stdin":"","call":"range_sums([-3, 8, -1], [(0, 2), (1, 3), (0, 3)])\n","expect":"[5, 7, 4]","hidden":true,"returns":true},{"name":"range_sums([4], [])","stdin":"","call":"range_sums([4], [])\n","expect":"[]","hidden":true,"returns":true},{"name":"sum(range_sums([i % 101 - 50 for i in range(200_000)], [(i % 1000, 199_000 + i % 1000) for i in range(100_000)]))","stdin":"","call":"sum(range_sums([i % 101 - 50 for i in range(200_000)], [(i % 1000, 199_000 + i % 1000) for i in range(100_000)]))\n","expect":"674500","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]},{"task":"Write `longest_stretch(values, k)`: the length of the longest stretch of values in a row that adds up to exactly `k`, or 0 if there is none. Values can be negative. The hidden test has 200,000 values, so trying every start and end is too slow.\n\nFor example, `longest_stretch([3, 4, -2, 5, 1, -4, 6], 7)` returns (3, 4, -2, 5, 1, -4):\n\n```\n6\n```\n","kind":"code","seed":"def longest_stretch(values, k):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(longest_stretch([3, 4, -2, 5, 1, -4, 6], 7))\n","solution":"def longest_stretch(values, k):\n    first = {0: 0}\n    total = 0\n    best = 0\n    for i, value in enumerate(values):\n        total += value\n        if total - k in first:\n            best = max(best, i + 1 - first[total - k])\n        if total not in first:\n            first[total] = i + 1\n    return best\n\n\nif __name__ == \"__main__\":\n    print(longest_stretch([3, 4, -2, 5, 1, -4, 6], 7))\n","hints":["As in `count_stretches`, a stretch adds up to `k` when an earlier prefix sum equals `total - k`. This time you want the earliest such prefix, as it gives the longest stretch.","Keep a dict that maps each prefix sum to the first position where it appeared, starting with `{0: 0}`. After adding `values[i]`, the prefix covers `i + 1` values, so a match at position `p` gives a stretch of `i + 1 - p` values.","Only store a prefix sum the first time you see it: a later, equal prefix sum would give shorter stretches."],"tests":[{"name":"longest_stretch([3, 4, -2, 5, 1, -4, 6], 7)","stdin":"","call":"longest_stretch([3, 4, -2, 5, 1, -4, 6], 7)\n","expect":"6","hidden":false,"returns":true},{"name":"longest_stretch([], 0)","stdin":"","call":"longest_stretch([], 0)\n","expect":"0","hidden":true,"returns":true},{"name":"longest_stretch([1, 2, 3], 10)","stdin":"","call":"longest_stretch([1, 2, 3], 10)\n","expect":"0","hidden":true,"returns":true},{"name":"longest_stretch([1, -1, 1, -1], 0)","stdin":"","call":"longest_stretch([1, -1, 1, -1], 0)\n","expect":"4","hidden":true,"returns":true},{"name":"longest_stretch([5, 1, 5], 5)","stdin":"","call":"longest_stretch([5, 1, 5], 5)\n","expect":"1","hidden":true,"returns":true},{"name":"longest_stretch([2, 3, -3, 4, 1], 4)","stdin":"","call":"longest_stretch([2, 3, -3, 4, 1], 4)\n","expect":"3","hidden":true,"returns":true},{"name":"longest_stretch([i % 7 - 2 for i in range(200_000)], 1_000)","stdin":"","call":"longest_stretch([i % 7 - 2 for i in range(200_000)], 1_000)\n","expect":"1003","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[]}]}