Back

Resumable Iterator with Multi-Dimensional Support

Low-Level DesignCodingOnsitePhoneMachine Learning EngineerSoftware EngineerReported Mar, 2026

Problem Overview

You need to implement a resumable iterator system that supports pause and resume functionality through state management. The problem progressively adds complexity through four parts: defining an abstract interface, implementing a basic resumable list iterator, building a 2D iterator, and extending to 3D iteration. The key challenge is maintaining state that can be serialized and restored to resume iteration from exact positions.

Part 1: Design the Abstract Interface

Define an abstract base class or interface for a resumable iterator that establishes the contract for state management and iteration.

Requirements

Define an abstract class ResumableIterator with core methods:

iter() and next() for Python iterator protocol

get_state() to capture current position

set_state(state) to restore position

State must be serializable (e.g., can be converted to/from dict or JSON)

Write unit tests for the interface expectations

Sample Interface Design

from abc import ABC, abstractmethod

from typing import Any, Dict

class ResumableIterator(ABC):

"""Abstract base class for iterators that support pause/resume"""

@abstractmethod

def iter(self):

"""Return the iterator object"""

return self

@abstractmethod

def next(self):

"""Return the next item or raise StopIteration"""

pass

@abstractmethod

def get_state(self) -> Dict[str, Any]:

"""

Capture the current state of the iterator.

Returns a serializable dictionary representing current position.

"""

pass

@abstractmethod

def set_state(self, state: Dict[str, Any]) -> None:

"""

Restore the iterator to a previously saved state.

Args:

state: Dictionary returned from a previous get_state() call

"""

pass

Example Unit Test Structure

import unittest

class TestResumableIterator(unittest.TestCase):
    def test_basic_iteration(self):
        """Test that iterator can iterate through all elements"""
        # Implementation specific
        pass

    def test_get_state_returns_serializable(self):
        """Test that get_state returns a JSON-serializable dict"""
        iterator = SomeResumableIterator([1, 2, 3])
        state = iterator.get_state()
        self.assertIsInstance(state, dict)
        # Verify it can be serialized
        import json
        json.dumps(state)  # Should not raise

    def test_set_state_restores_position(self):
        """Test that set_state correctly restores iteration position"""
        iterator = SomeResumableIterator([1, 2, 3])
        next(iterator)  # Advance to position 1
        state = iterator.get_state()
        next(iterator)  # Advance to position 2
        iterator.set_state(state)  # Restore to position 1
        self.assertEqual(next(iterator), 2)  # Should get element at index 1

    def test_pause_and_resume(self):
        """Test complete pause/resume cycle"""
        iterator = SomeResumableIterator([10, 20, 30, 40])
        self.assertEqual(next(iterator), 10)
        self.assertEqual(next(iterator), 20)

        # Pause
        state = iterator.get_state()

        # Create new iterator and resume
        new_iterator = SomeResumableIterator([10, 20, 30, 40])
        new_iterator.set_state(state)

        # Should continue from position 2
        self.assertEqual(next(new_iterator), 30)
        self.assertEqual(next(new_iterator), 40)

Part 2: Implement Basic Resumable List Iterator

Implement a concrete ResumableListIterator class that iterates through a list with pause/resume support.

Requirements

Inherit from the abstract interface defined in Part 1

Initialize with a list of items

Track current position (index) as state

Support pausing at any point and resuming from that exact position

Handle edge cases: empty lists, iteration beyond bounds

Example Implementation

class ResumableListIterator(ResumableIterator):
    def __init__(self, items: list):
        self.items = items
        self.index = 0

    def __iter__(self):
        return self

    def __next__(self):
        if self.index >= len(self.items):
            raise StopIteration

        result = self.items[self.index]
        self.index += 1
        return result

    def get_state(self) -> Dict[str, Any]:
        """Return current position and metadata"""
        return {
            'index': self.index,
            'total_items': len(self.items)
        }

    def set_state(self, state: Dict[str, Any]) -> None:
        """Restore to saved position"""
        self.index = state['index']
        # Optionally validate total_items matches
        if state.get('total_items') != len(self.items):
            raise ValueError("State was saved for different data")

Usage Example

# Create iterator
iterator = ResumableListIterator(['a', 'b', 'c', 'd', 'e'])

# Iterate partially
print(next(iterator))  # 'a'
print(next(iterator))  # 'b'

# Capture state
state = iterator.get_state()
print(state)  # {'index': 2, 'total_items': 5}

# Continue iterating
print(next(iterator))  # 'c'

# Restore to previous state
iterator.set_state(state)
print(next(iterator))  # 'c' again

# Continue to completion
print(next(iterator))  # 'd'
print(next(iterator))  # 'e'
# next(iterator) would raise StopIteration

Edge Cases to Consider

Empty list initialization

Setting state beyond list bounds

Setting state from incompatible iterator (different data)

Multiple get_state calls without advancing

Restoring state after iterator exhaustion

Part 3: 2D Resumable List Iterator

Build a 2D iterator that iterates through a list of lists (or a matrix) by extending the resumable-iterator idea from Part 2 to multiple dimensions. This requires tracking position across multiple dimensions.

Requirements

Accept a 2D list (list of lists) as input

Flatten the iteration: iterate through all elements row by row

Track both outer index (which inner list) and inner index (position within that list)

Handle edge cases:

Empty outer list

Empty inner lists (skip them)

Irregular nested lists (different lengths)

Handle skipping empty inner lists in next(), either iteratively or recursively

Key Challenge: Corner Cases

The most common issues arise from:

Empty inner lists that need to be skipped

Moving from one inner list to the next

Detecting when all lists are exhausted

Properly handling state restoration when positioned between lists

Example Implementation

class Resumable2DIterator(ResumableIterator):
    def __init__(self, items: list[list]):
        self.items = items
        self.outer_index = 0
        self.inner_index = 0

    def __iter__(self):
        return self

    def __next__(self):
        # Skip empty inner lists (recursion handles this elegantly)
        while self.outer_index < len(self.items):
            current_list = self.items[self.outer_index]

            # If current inner list is empty, move to next
            if len(current_list) == 0:
                self.outer_index += 1
                self.inner_index = 0
                continue

            # If we've exhausted current inner list, move to next
            if self.inner_index >= len(current_list):
                self.outer_index += 1
                self.inner_index = 0
                continue

            # Return current element and advance inner index
            result = current_list[self.inner_index]
            self.inner_index += 1
            return result

        # All lists exhausted
        raise StopIteration

    def get_state(self) -> Dict[str, Any]:
        return {
            'outer_index': self.outer_index,
            'inner_index': self.inner_index
        }

    def set_state(self, state: Dict[str, Any]) -> None:
        self.outer_index = state['outer_index']
        self.inner_index = state['inner_index']

Alternative Recursive Approach

def next(self):

"""More explicit recursive structure"""

if self.outer_index >= len(self.items):

raise StopIteration

current_list = self.items[self.outer_index]

If current list is exhausted or empty, move to next

if self.inner_index >= len(current_list):

self.outer_index += 1

self.inner_index = 0

return self.next() # Recursive call to try next list

Return current element

result = current_list[self.inner_index]

self.inner_index += 1

return result

Usage Example

data = [
    [1, 2, 3],
    [],           # Empty inner list
    [4, 5],
    [6]
]

iterator = Resumable2DIterator(data)

# Iterate partially
print(next(iterator))  # 1
print(next(iterator))  # 2
print(next(iterator))  # 3

# Capture state (finished first list, will move to empty list next)
state = iterator.get_state()

# Continue
print(next(iterator))  # 4 (skips empty list)
print(next(iterator))  # 5

# Restore state
iterator.set_state(state)
print(next(iterator))  # 4 (correctly resumes, skipping empty list)
print(next(iterator))  # 5
print(next(iterator))  # 6
# next(iterator) raises StopIteration

Critical Corner Cases

Test Case 1: All empty inner lists

data = [[], [], []]

iterator = Resumable2DIterator(data)

Should immediately raise StopIteration

Test Case 2: Empty at boundaries

data = [[], [1, 2], []]

iterator = Resumable2DIterator(data)

assert list(iterator) == [1, 2]

Test Case 3: State capture at list boundary

data = [[1], [2], [3]]

iterator = Resumable2DIterator(data)

next(iterator) # 1

state = iterator.get_state() # outer_index=0, inner_index=1

next(iterator) # 2

iterator.set_state(state)

assert next(iterator) == 2 # Should correctly resume

Test Case 4: Irregular lengths

data = [[1], [2, 3, 4, 5], [6, 7]]

iterator = Resumable2DIterator(data)

state_after_3 = None

for i, val in enumerate(iterator):

if i == 2: # After getting 3

state_after_3 = iterator.get_state()

iterator.set_state(state_after_3)

assert next(iterator) == 4

Part 4: 3D Resumable Iterator (Advanced)

Extend the concept to three dimensions, iterating through a list of lists of lists.

Requirements

Accept a 3D list structure

Track three indices: outer, middle, and inner

Skip empty lists at any level

Properly handle state across all three dimensions

Manage complexity of nested empty list handling

Implementation Approach

class Resumable3DIterator(ResumableIterator):
    def __init__(self, items: list[list[list]]):
        self.items = items
        self.outer_index = 0
        self.middle_index = 0
        self.inner_index = 0

    def __iter__(self):
        return self

    def __next__(self):
        while self.outer_index < len(self.items):
            if self.middle_index >= len(self.items[self.outer_index]):
                # Exhausted middle dimension, move to next outer
                self.outer_index += 1
                self.middle_index = 0
                self.inner_index = 0
                continue

            current_middle = self.items[self.outer_index][self.middle_index]

            if len(current_middle) == 0:
                # Empty inner list, skip
                self.middle_index += 1
                self.inner_index = 0
                continue

            if self.inner_index >= len(current_middle):
                # Exhausted inner dimension, move to next middle
                self.middle_index += 1
                self.inner_index = 0
                continue

            # Return current element
            result = current_middle[self.inner_index]
            self.inner_index += 1
            return result

        raise StopIteration

    def get_state(self) -> Dict[str, Any]:
        return {
            'outer_index': self.outer_index,
            'middle_index': self.middle_index,
            'inner_index': self.inner_index
        }

    def set_state(self, state: Dict[str, Any]) -> None:
        self.outer_index = state['outer_index']
        self.middle_index = state['middle_index']
        self.inner_index = state['inner_index']

Usage Example

data = [
    [
        [1, 2],
        [3]
    ],
    [
        [],
        [4, 5, 6]
    ],
    [
        [7]
    ]
]

iterator = Resumable3DIterator(data)

# Iterate and pause
print(next(iterator))  # 1
print(next(iterator))  # 2
print(next(iterator))  # 3

state = iterator.get_state()

print(next(iterator))  # 4 (skips empty list)

# Restore
iterator.set_state(state)
print(next(iterator))  # 4
print(next(iterator))  # 5

Discussion Points

If you only have 5 minutes (as mentioned in the interview reports), focus on:

State structure: three indices needed

Nested while/if logic to handle three dimensions

Skip logic for empty lists at each level

How get_state() and set_state() would work

Pseudocode approach for explanation:

def next(self):

Loop through outer dimension

while not at_end_of_outer:

Loop through middle dimension

while not at_end_of_middle:

Check if inner list is empty or exhausted

if inner_list_has_elements:

return element and advance inner_index

else:

advance to next middle_index

advance to next outer_index

raise StopIteration

Async/Concurrent Extension (Bonus Concept)

Some interview variations may ask about asynchronous iteration for file I/O or network sources.

Key Concepts

Use async def anext() instead of next()

Support async for iteration

State must capture async operation progress (byte offset, not just line number)

File handles can't be serialized, only their position

Example Interface

class AsyncResumableIterator(ABC):
    @abstractmethod
    async def __anext__(self):
        """Async version of __next__"""
        pass

    @abstractmethod
    def get_state(self) -> Dict[str, Any]:
        """State must not include file handles, only positions"""
        pass

    @abstractmethod
    async def set_state(self, state: Dict[str, Any]) -> None:
        """May need async if reopening files"""
        pass

# Usage
async for item in async_iterator:
    if some_condition:
        state = async_iterator.get_state()
        # Save state to database or file
        break

# Later, restore
await async_iterator.set_state(saved_state)
async for item in async_iterator:
    # Continues from saved position
    process(item)

Testing Strategy

def test_comprehensive_2d():

"""Test all edge cases for 2D iterator"""

test_cases = [

(input, expected_output)

([[1, 2], [3, 4]], [1, 2, 3, 4]),

([[], [1], []], [1]),

([[], [], []], []),

([[1]], [1]),

([[1, 2, 3]], [1, 2, 3]),

([[1], [2], [3]], [1, 2, 3]),

]

for input_data, expected in test_cases:

iterator = Resumable2DIterator(input_data)

result = list(iterator)

assert result == expected, f"Failed for {input_data}"

Test pause/resume

iterator = Resumable2DIterator([[1, 2], [3, 4]])

next(iterator)

next(iterator)

state = iterator.get_state()

next(iterator)

iterator.set_state(state)

assert next(iterator) == 3

Time and Space Complexity

Time Complexity:

next(): O(1) amortized per element (may skip multiple empty lists, but each list is skipped at most once)

Complete iteration of N elements: O(N)

get_state(): O(1)

set_state(): O(1)

Space Complexity:

O(1) - only storing index positions, not copying data

State dictionary: O(1) - constant number of fields regardless of data size

Real-World Applications

This pattern is useful for:

ETL Pipelines: Processing large datasets with checkpointing

Distributed Processing: Pausing work on one machine, resuming on another

Long-Running Jobs: Surviving process restarts or crashes

Rate-Limited APIs: Pausing when hitting limits, resuming later

Streaming Data: Buffering and resuming streams across network interruptions


Auto-save enabled
Loading editor…
Output
Run your code to see the output here.