https://leetcode.com/problems/zigzag-iterator/
Zigzag Iterator
Given two 1d vectors, implement an iterator to return their elements alternately.
For example, given two 1d vectors:
v1 = [1, 2] v2 = [3, 4, 5, 6]
By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1, 3, 2, 4, 5, 6].
Follow up: What if you are given k 1d vectors? How well can your code be extended to such cases?
Clarification for the follow up question – Update (2015-09-18):
The “Zigzag” order is not clearly defined and is ambiguous for k > 2 cases. If “Zigzag” does not look right to you, replace “Zigzag” with “Cyclic”. For example, given the following input:
[1,2,3] [4,5,6,7] [8,9]
It should return [1,4,8,2,5,9,3,6,7].
My verbose code:
class ZigzagIterator(object):
def __init__(self, v1, v2):
"""
Initialize your data structure here.
:type v1: List[int]
:type v2: List[int]
"""
self.v1 = v1
self.v2 = v2
self.pointer = 0
if len(v1) != 0:
self.l = 0
else:
self.l = 1
self.x = 0
def next(self):
"""
:rtype: int
"""
n = None
if self.l == 0:
n = self.v1[self.x]
if self.x < len(self.v2):
self.l = 1
# keep self.x same
else:
# keep self.l same
self.x += 1
else:
n = self.v2[self.x]
if self.x < len(self.v1) - 1:
self.l = 0
self.x += 1
else:
self.x += 1
self.pointer += 1
return n
def hasNext(self):
"""
:rtype: bool
"""
return self.pointer < len(self.v1)+len(self.v2)
# Your ZigzagIterator object will be instantiated and called as such:
# i, v = ZigzagIterator(v1, v2), []
# while i.hasNext(): v.append(i.next())
However, there are more elegant codes which leverage Python generator (https://leetcode.com/discuss/57997/python-o-1-space-solutions):
class ZigzagIterator(object):
def __init__(self, v1, v2):
self.vals = (v[i] for i in itertools.count() for v in (v1, v2) if i < len(v))
self.n = len(v1) + len(v2)
def next(self):
self.n -= 1
return next(self.vals)
def hasNext(self):
return self.n > 0
Here, self.vals = (v[i] for i in itertools.count() for v in (v1, v2) if i < len(v)) uses `itertools.count()` as a generator to generate infinite integer `i`. For each `i`, if `i < len(v)` then the list `self.vals` will add v[i]. This essentially constitutes a zigzag output in the initialization. However, since `itertools.count()` is just a generator, `self.vals` don’t occupy the memory: every element in `self.val` will be yielded whenever you call `next(self.vals)`.