Medium
Flatten 2D Vector — Python
Full explanation · Time O(1) · Space O(1)
# Time: O(1)
# Space: O(1)
from collections import deque
class Vector2D(object):
def __init__(self, vec2d):
"""
Initialize your data structure here.
:type vec2d: List[List[int]]
"""
self.stack = deque((len(v), iter(v)) for v in vec2d if v)
def next(self):
"""
:rtype: int
"""
length, iterator = self.stack.popleft()
if length > 1:
self.stack.appendleft((length-1, iterator))
return next(iterator)
def hasNext(self):
"""
:rtype: bool
"""
return bool(self.stack)