Approach 1: Keep track of list number, item number pair
class ZigzagIterator(object):
def __init__(self, v1, v2):
"""
Initialize your data structure here.
:type v1: List[int]
:type v2: List[int]
"""
self.curr = 0
self.idx = [0, 0]
self.arr = [v1, v2]
self.size = 2
def next(self):
"""
:rtype: int
"""
if self.idx[self.curr] == len(self.arr[self.curr]):
del self.arr[self.curr]
del self.idx[self.curr]
self.size -= 1
self.curr = (self.curr) % self.size
ret = self.arr[self.curr][self.idx[self.curr]]
self.idx[self.curr] += 1
self.curr = (self.curr + 1) % self.size
return ret
def hasNext(self):
"""
:rtype: bool
"""
for i in range(self.size):
if self.idx[i] < len(self.arr[i]):
return True
return False
# Your ZigzagIterator object will be instantiated and called as such:
# i, v = ZigzagIterator(v1, v2), []
# while i.hasNext(): v.append(i.next())
Storing the index number of current element in each list allows us to iterate over each of the lists to return the elements one by one.
Time Complexity: O(N), where N = length(v1) + length(v2)
Space Complexity: O(N), where N = length(v1) + length(v2)
