Skip to content

← Stdlib Mastery step 11 of 55

Easy Primitives

itertools: take n across chunks without over-consuming

Take the first n items across a stream of chunks, starting no chunk you do not need.

Implement:

def take_flattened(chunks: Iterable[Iterable[int]], n: int) -> list[int]: ...

def solve(chunks: list[list[int]], n: int) -> tuple[list[int], int]: ...

solve wraps chunks in a counting generator that increments a counter each time it hands over a chunk, passes that generator to take_flattened, and returns (items, chunks_started).

The counter is the assertion. A solution that reads every chunk before slicing gets the items right and the count wrong, and there is no way to fake the count without actually understanding the consumption order:

  • n = 4 over chunks of size 2 must start exactly 2 chunks, not 3. islice checks its counter before pulling, so it never asks for the item that would trigger the third chunk.
  • n = 0 must start 0 chunks — creating the pipeline consumes nothing.
  • n larger than the total must start every chunk (the last pull is what discovers the stream has ended).
  • Empty chunks are still started: [[], [], [7]] with n = 1 starts 3.

The production consequence. list(a) + list(b) reads both sides into memory purely in order to walk them once. On a 10 GB log that is the difference between a job that works in 30 MB and one that gets OOM-killed. The same applies one level up: chain(*sources) unpacks the outer iterable and so consumes it entirely before yielding anything, while chain.from_iterable(sources) pulls the next source only when the current one runs out. If sources is a generator of open file handles, those two lines differ by “how many files are open at once”.

Your submission must pass mypy --strict. Take Iterable[Iterable[int]], not list[list[int]] — the whole point is that the source need not be a sequence.

Returns a tuple, not a list. The harness compares container types exactly.

Loading visualization…