BinarySearch 0732 Two-Dimensional List Iterator

Problem statement

https://binarysearch.com/problems/Two-Dimensional-List-Iterator/

Solution

Equal to Leetcode 0251. Flatten 2D Vector.

Complexity

See complexities of leetcode problem.

Code

class TwoDimensionalIterator:
    def __init__(self, vec):
        self.vec = vec
        self.x = 0
        self.y = 0
        
    def helper(self):
        while self.x < len(self.vec) and self.y == len(self.vec[self.x]):
            self.x += 1
            self.y = 0

    def next(self):
        self.helper()
        ans = self.vec[self.x][self.y]
        self.y += 1
        return ans

    def hasnext(self):
        self.helper()
        return self.x < len(self.vec)