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)`.