Skip to content

← Memory, Copies and Slow Code step 32 of 38

Medium End-to-End

Speed Up a Moving Average with Prefix Sums

Write a moving average that takes the same time whatever the window size, and compare it with the direct version. Use Python lists. You don't need NumPy.

def moving_average_naive(xs: Sequence[float], window: int) -> list[float]: ...
def moving_average_cumsum(xs: Sequence[float], window: int) -> list[float]: ...
def choose(n: int, vectorised: bool) -> str: ...

Rules for both averages

  • Return one average per full window, in input order. There are max(len(xs) - window + 1, 0) of them, so a window longer than the input gives [].
  • Raise ValueError when window is less than 1, even for empty input.
  • Don't change xs.

The starter already includes moving_average_naive. Keep it as the version to compare against.

Write the prefix-sum version

moving_average_cumsum must visit each value a fixed number of times, however wide the window is:

  1. Build a prefix list that starts with 0.0, followed by the running totals of xs from left to right.
  2. The sum of the window xs[i:i + window] is prefix[i + window] - prefix[i].
  3. Divide that sum by window.
assert moving_average_cumsum([1.0, 2.0, 3.0, 4.0, 5.0], 3) == [2.0, 3.0, 4.0]

The starter leaves out the leading 0.0 and subtracts the wrong prefix, so every window misses its first value. A window of 1 shows the bug at once, because it should return the input unchanged.

The two versions add the same numbers in a different order, so their float results can differ slightly. The driver compares them with a tolerance, not with ==.

A fixed rule for choose

choose returns a label by a rule that exists only for this exercise:

  • Return "list" when vectorised is False.
  • Otherwise return "list" when n < 500 and "numpy" when n >= 500.

The 500 is a made-up threshold, not a measured point where NumPy gets faster. The starter ignores vectorised.

What the driver reports

Keep solve as it is.

  • If generate > 0, it replaces xs with [((i * 48271) % 10007) / 10007 for i in range(generate)].
  • method is "naive", "cumsum" or "both". With "both", it runs both and reports the prefix-sum results.
  • agree is True when both ran and every pair of results matches within math.isclose(a, b, rel_tol=1e-9, abs_tol=1e-12). With one method it is always True, though the driver compared nothing.
  • count is the number of averages. head and tail are the first four and the last four, each formatted with six decimal places. A short result can appear whole in both.
  • choice is choose(choose_n, vectorised).

The largest test runs "cumsum" on 200,000 values with a window of 5,000. The direct version would add up about 975 million values there. Your submission must pass mypy --strict.