# The Custom Iterator

> Some sequences follow their own rules.

Canonical URL: <https://datadriven.io/problems/the_custom_iterator>

Domain: Python · Difficulty: medium · Seniority: L4

## Problem

Implement a Range class accepting (start, stop, step=1). Support __iter__ and __next__. The test harness constructs Range(*args) then calls 'iterate' which collects list(iter(range_instance)). Return None for constructor calls and the collected list for 'iterate' calls.

## Worked solution and explanation

### Why this problem exists in real interviews

This tests a deeper understanding of the **iterator protocol** by reimplementing `range()` as a class. It probes mastery of `__iter__`, `__next__`, `StopIteration`, and handling of optional parameters with default step values.

---

### Break down the requirements

#### Step 1: Handle constructor arguments

Support (stop), (start, stop), and (start, stop, step) signatures. Default start to 0, step to 1.

#### Step 2: Implement __iter__ to return self

The Range object acts as its own iterator.

#### Step 3: Implement __next__ with step advancement

Track the current value, check against stop, yield, and advance by step. Raise `StopIteration` when done.

---

### The solution

**Class-based range with iterator protocol**

```python
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):
        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
```

> **Time and Space Complexity**
>
> **Time:** O(1) per `__next__` call. O((stop - start) / step) for full iteration.
> 
> **Space:** O(1). Only scalar state variables.

> **Interviewers Watch For**
>
> Whether you handle negative step values correctly. `Range(10, 0, -2)` should yield `[10, 8, 6, 4, 2]`.

> **Common Pitfall**
>
> Forgetting to handle the case where `step` is negative. The stop condition flips: you stop when `current` is less than or equal to `stop`.

---

## Common follow-up questions

- How would you make Range reusable across multiple for loops? _(Tests returning a new iterator from `__iter__` instead of self.)_
- What happens if step is 0? _(Tests error handling: should raise `ValueError` like built-in `range`.)_
- How does Python's built-in range achieve O(1) membership testing? _(Tests knowledge of `__contains__` and arithmetic bounds checking.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/the_custom_iterator)
- [Python Interview Questions](https://datadriven.io/python-interview-questions)
- [Data Engineering Interview Prep Guide](https://datadriven.io/data-engineer-interview-prep)
- [Daily Challenge](https://datadriven.io/daily)

---

Source: DataDriven (https://datadriven.io). 100% free data engineering interview prep. Live code execution against Postgres 16, Python 3.11, and Spark sandboxes. No paywall, no premium tier, no signup gate.