Leetcode 0251. Flatten 2D Vector

Problem statement

https://leetcode.com/problems/flatten-2d-vector/

Solution

Keep 3 pieces of information: self.vec is our list of lists, self.x is index in this list and self.y is index in self.vec[self.x]. Also we use helper function which will keep invariant: it will always point to existing index, that is if elements ended in some list, it will point to the next non-empty one.

1. next(self) now is just calling for helper, save element, increase self.y and return saved element.
2. hasNext(self) is just calling for helper and checking that we are not in the very end, that is self.x < len(self.vec).

Complexity

Imagine, that we have $n$ elements in total in list of lists and $m$ elements in total if we count empty lists. Then to construct we need $O(1)$ time and to run next and hasNext it is in average $O(m/n)$, that is to traverse all list with $n$ elements we need $O(m)$ operations. Space complexity is $O(1)$.

Code

class Vector2D:
    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)