Write a class that implements the iterator protocol to count down from n to 1.
Define __iter__ returning self and __next__ returning the next value, raising StopIteration when the countdown ends. Store the current position on the instance. Each next call is O(1) and the class uses O(1) space.
Explain it simply
An iterator needs two methods. __iter__ returns the iterator itself, which makes it usable in a for loop. __next__ returns the next value and signals the end by raising StopIteration.
Line by line
class CountDown:
def __init__(self, start):
self.current = start
def __iter__(self):
return self
def __next__(self):
if self.current <= 0:
raise StopIteration
value = self.current
self.current -= 1
return value
for n in CountDown(3):
print(n)
# Output: 3
# Output: 2
# Output: 1
print(list(CountDown(5))) # Output: [5, 4, 3, 2, 1]
print(sum(CountDown(4))) # Output: 10
it = CountDown(1)
print(next(it)) # Output: 1
print(list(it)) # Output: []The state lives in self.current, so the object remembers where it is. Once it reaches zero, it stays exhausted.
Complexity
Time O(1) per value, O(n) for the full countdown. Space O(1), since no list of numbers is stored.
What to say in the interview
"I implement __iter__ returning self and __next__ that returns the next value or raises StopIteration. The iterator keeps its own position, so it is single-use. In practice a generator function with yield does the same thing in three lines, but the class shows how the protocol works."