# The Counting Machine

> It knows where it stopped last time.

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

Domain: Python · Difficulty: medium · Seniority: L4

## Problem

Implement CountUp(values). Its iterator yields values[i] * (i + 1) for i = 0..len(values)-1. The test harness constructs the object then collects list(iter(obj)) and expects the produced list.

## Worked solution and explanation

### Why this problem exists in real interviews

This tests understanding of the **iterator protocol** in Python: `__iter__` and `__next__`. It probes whether a candidate can implement a stateful class that tracks position and produces transformed values on demand.

---

### Break down the requirements

#### Step 1: Store the input list and initialize position to 0

The constructor saves the data; iteration state starts at the beginning.

#### Step 2: Implement __iter__ returning self

This makes the object itself the iterator.

#### Step 3: Implement __next__ with position tracking

Multiply the current element by its 1-based index, advance the pointer, and raise `StopIteration` when exhausted.

---

### The solution

**Stateful iterator with index-based multiplication**

```python
class CountUp:
    def __init__(self, data: list):
        self.data = data
        self.index = 0
    def __iter__(self):
        return self
    def __next__(self):
        if self.index >= len(self.data):
            raise StopIteration
        value = self.data[self.index] * (self.index + 1)
        self.index += 1
        return value
```

> **Time and Space Complexity**
>
> **Time:** O(1) per `__next__` call. O(n) to iterate the full sequence.
> 
> **Space:** O(1) beyond the stored input. Only the index counter is extra state.

> **Interviewers Watch For**
>
> Whether you raise `StopIteration` correctly. Returning `None` instead of raising the exception breaks `for` loop integration.

> **Common Pitfall**
>
> Using 0-based indexing for the multiplier. The problem specifies 1-based position, so the first element is multiplied by 1, not 0.

---

## Common follow-up questions

- How would you make this iterator reusable (resettable)? _(Tests adding a reset method or returning a new iterator from `__iter__`.)_
- What is the difference between an iterator and an iterable? _(Tests conceptual understanding: iterables have `__iter__`, iterators have both `__iter__` and `__next__`.)_
- How would you implement this as a generator instead? _(Tests `yield` syntax as a simpler alternative to the class-based protocol.)_

## Related

- [All practice problems](https://datadriven.io/problems)
- [Mock interview mode](https://datadriven.io/interview/the_counting_machine)
- [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.