The Custom Iterator

A medium Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
medium
Seniority
L4
Asked in
Foursquare

Problem

Our Range objects get walked more than once, so iterating the same instance again must yield its full sequence again, and a negative step counts down the way the built-in range does. Build Range as its own iterator, taking the built-in's one, two and three argument forms. Then have run_operations replay operations against the matching arguments: "Range" builds a new instance and records None, and "iterate" records every value the most recent instance yields.

Some sequences follow their own rules.

Example 1

Input:

arguments = [[0,5,1],[]], operations = ["Range","iterate"]

Output:

[null,[0,1,2,3,4]]

Example 2

Input:

arguments = [[0,10,2],[]], operations = ["Range","iterate"]

Output:

[null,[0,2,4,6,8]]

Worked solution and explanation

The trap

This is a re-iterable sequence wearing the costume of range(). Almost everyone gets __next__ right: yield, advance, raise StopIteration at the boundary. What separates candidates is where the cursor gets reset. Set self.current = self.start only in __init__ and have __iter__ just return self, and the object works exactly once. The second 'iterate' on the same instance gets an already drained iterator and you record [] where [0, 1, 2] belongs.

Building it

Step 1: Normalize the constructor once

*args with one, two or three values maps onto start, stop and step the way the built-in does: a single value is the stop, start defaults to 0 and step to 1. Decide this in __init__ so __next__ never cares which form the caller used.

Step 2: Rewind the cursor in __iter__

__iter__ is the moment a new pass begins, so it is where self.current goes back to self.start before returning self. That single line is the difference between an iterable and a single-use iterator.

Step 3: Let the step sign pick the boundary

Counting up stops once current >= stop; counting down stops once current <= stop. A single ascending comparison silently returns an empty countdown, because 10 already looks past 0 from below.

Step 4: Replay the operations

run_operations holds the most recently built Range. 'Range' replaces it and records None; 'iterate' records list(current), which calls __iter__ first and therefore starts a fresh pass every time.

A re-iterable `Range`
class Range:
    def __init__(self, *args):
        if len(args) == 1:
            self.start = 0
            self.stop = args[0]
            self.step = 1
        elif len(args) == 2:
            self.start = args[0]
            self.stop = args[1]
            self.step = 1
        else:
            self.start = args[0]
            self.stop = args[1]
            self.step = args[2]
        self.current = self.start

    def __iter__(self):
        # Every new pass starts from the beginning.
        self.current = self.start
        return self

    def __next__(self):
        if self.step > 0 and self.current >= self.stop:
            raise StopIteration
        if self.step < 0 and self.current <= self.stop:
            raise StopIteration
        value = self.current
        self.current += self.step
        return value


def run_operations(operations, arguments):
    results = []
    current = None
    for op, args in zip(operations, arguments):
        if op == "Range":
            current = Range(*args)
            results.append(None)
        elif op == "iterate":
            results.append(list(iter(current)))
        else:
            raise ValueError(f"unknown operation: {op}")
    return results
Single-use
Cursor set only in __init__, __iter__ returns self. The first pass yields [0, 1, 2]; the second yields [].
Re-iterable
Cursor rewound in __iter__. Every pass over the same instance yields [0, 1, 2].

Common follow-up questions