Skip to content

← Under the Hood: Objects, Memory, Speed step 34 of 35

Medium End-to-End

Moving average twice: naive, cumulative-sum, and knowing which to use

Implement a rolling mean two ways, prove they agree, and encode the list-versus-NumPy decision as a function rather than a reflex.

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: ...

Both averages return one value per full window — len(xs) - window + 1 results, or [] when the window is longer than the input — and both raise ValueError for window < 1.

moving_average_naive sums each window independently. O(n × window). Provided and correct; it is the reference the cumulative version is checked against.

moving_average_cumsum builds a prefix-sum array once and takes differences. O(n). The starter has the classic off-by-one: it accumulates without a leading 0.0, so prefix[i + window - 1] - prefix[i] is short by one term in every window and shifted. The output has the right length, so a len() assertion passes and “it runs” passes — which is exactly why this bug reaches production.

choose(n, vectorised) returns "list" or "numpy" under two rules:

  • If you are going to loop in Python over the elements, the answer is "list" at any n. Indexing a NumPy array from Python boxes every element into a fresh scalar object, so you pay the array’s cost and then re-box anyway — strictly worse than the list you started with.
  • Otherwise, "list" below 500 elements, where NumPy’s per-call dispatch dominates the arithmetic, and "numpy" at or above it.

The starter ignores vectorised entirely, which is the “big data means NumPy” reflex — it drops the only fact that actually decides the question.

solve(xs, window, generate, method, choose_n, vectorised) generates the input when generate > 0, runs "naive", "cumsum" or "both", and returns the result count, the first and last four values as fixed-point strings, whether the two methods agreed (for "both"), and choose‘s answer.

One hidden case runs 200,000 values with a 5,000-wide window, in "cumsum" mode. That is 10⁹ operations naively and 2 × 10⁵ with prefix sums — the algorithm is the difference between 0.02 seconds and several minutes, and no amount of vectorising the naive version closes it.

Your submission must pass mypy --strict.

Loading visualization…