251. Flatten 2D Vector

On LeetCode ->

class Vector2D:
    def __init__(self, vec: list[list[int]]):
        pass

    def next(self):
        pass

    def hasNext(self):
        pass

Problem¶

Implement a lazy iterator over a list of integer lists, visiting elements row by row.

next() returns and consumes the next integer; calls are guaranteed valid.
hasNext() checks availability without consuming an integer.
Empty rows must be skipped.

vec = [[], [1, 2], [], [3], []]
calls:   hasNext(), hasNext(), next(), next(), hasNext(), next(), hasNext()
returns: True,      True,      1,      2,      True,      3,      False

Key trick¶

Track a row and column; before either operation, skip exhausted rows, including empty ones.

Trap¶

Repeated hasNext() calls must not consume values, and next() must work without a preceding hasNext() call.

Why is it interesting?¶

It tests lazy traversal and iterator state management without copying the flattened data.

Python solution¶

class Vector2D:
    def __init__(self, vec: list[list[int]]):
        self.vec = vec
        self.r = 0
        self.c = 0

    def hasNext(self) -> bool:
        # Normalize the cursor to the next available value.
        while self.r < len(self.vec) and self.c == len(self.vec[self.r]):
            self.r += 1
            self.c = 0
        return self.r < len(self.vec)

    def next(self) -> int:
        # Also handle callers that did not first call hasNext().
        if not self.hasNext():
            raise StopIteration

        val = self.vec[self.r][self.c]
        self.c += 1
        return val

Time: Initialization is \(O(1)\). Either operation takes \(O(m)\) worst case when skipping rows; \(q\) operations take \(O(m + q)\) total, where \(m\) is the number of rows.

Auxiliary space: \(O(1)\); the input is referenced, not copied.

Comment on my solution¶

Your lookahead idea is valid, but the implementation mixes an iterator with a stack.

self.next = 0 shadows the next() method, so calling it fails; use self.val for a cached value.
if not self.hasNext checks a bound method, not availability; call self.hasNext().
iter(vec) has neither append() nor pop(), and remains truthy after exhaustion.
A stack of iterators fits arbitrarily nested lists; this fixed two-dimensional problem needs only two indices or an outer and inner iterator.

# Ideas:
# - It reminds me of "Flatten nested list iterator" problem

# WRONG
class Vector2D:
    def __init__(self, vec: list[list[int]]):
        self.stack = iter(vec)
        self.ready = False
        self.next = 0

    def next(self):
        if not self.hasNext:
            raise StopIteration
        self.ready = False
        return self.next

    def hasNext(self):
        if self.ready:
            return True

        while self.stack:
            try:
                item = next(self.stack)
            except StopIteration:
                self.stack.pop()
                continue

            if isinstance(item, int):
                self.ready = True
                self.next = item
                return True

            self.stack.append(iter(item))

        return False

# WRONG
# updated after reading AI comments
class Vector2D:
    def __init__(self, vec: list[list[int]]):
        self.stack = [iter(vec)]
        self.ready = False
        self.next = 0

    def next(self):
        if not self.hasNext():
            raise StopIteration
        self.ready = False
        return self.next

    def hasNext(self):
        if self.ready:
            return True

        while self.stack:
            try:
                item = next(self.stack[-1])
            except StopIteration:
                self.stack.pop()
                continue

            if isinstance(item, int):
                self.ready = True
                self.next = item
                return True

            self.stack.append(iter(item))

        return False

# WORKS
# Fixed by AI
# my error: `self.next = 0` shadows your `next()` method:
# `vector.next()` tries to call an integer.  Rename the cached value
# to `self.val`; the rest of your iterator-stack approach works.
class Vector2D:
    def __init__(self, vec: list[list[int]]):
        self.stack = [iter(vec)]
        self.ready = False
        self.val = 0

    def next(self) -> int:
        if not self.hasNext():
            raise StopIteration
        self.ready = False
        return self.val

    def hasNext(self) -> bool:
        # Repeated checks must not advance past the cached value.
        if self.ready:
            return True

        while self.stack:
            try:
                item = next(self.stack[-1])
            except StopIteration:
                self.stack.pop()
                continue

            if isinstance(item, int):
                self.val = item
                self.ready = True
                return True

            self.stack.append(iter(item))

        return False

Extra¶

Problem description¶

Design an iterator to flatten a 2D vector. It should support the next and hasNext operations.

Implement the Vector2D class:

Vector2D(int[][] vec) initializes the object with the 2D vector vec.
next() returns the next element from the 2D vector and moves the pointer one step forward. You may assume that all the calls to next are valid.
hasNext() returns true if there are still some elements in the vector, and false otherwise.

Example 1:

Input
["Vector2D", "next", "next", "next", "hasNext", "hasNext", "next", "hasNext"]
[[[[1, 2], [3], [4]]], [], [], [], [], [], [], []]
Output
[null, 1, 2, 3, true, true, 4, false]

Explanation
Vector2D vector2D = new Vector2D([[1, 2], [3], [4]]);
vector2D.next();    // return 1
vector2D.next();    // return 2
vector2D.next();    // return 3
vector2D.hasNext(); // return True
vector2D.hasNext(); // return True
vector2D.next();    // return 4
vector2D.hasNext(); // return False

Constraints:

0 <= vec.length <= 200
0 <= vec[i].length <= 500
-500 <= vec[i][j] <= 500
At most 10^5 calls will be made to next and hasNext.

Follow up: As an added challenge, try to code it using only iterators in C++ or iterators in Java.