Zigzag Iterator

Problem: Zigzag Iterator

From the sample, we can see clearly that the zigzag iterator we want is like index of a matrix. What's different is that the iterator can deal with rows of different length, which means we need to modify column index.

Code in Python:

class ZigzagIterator(object):

    def __init__(self, v1, v2):
        """
        Initialize your data structure here.
        :type v1: List[int]
        :type v2: List[int]
        """
        self.v = [v1[:], v2[:]]
        self.i, self.j = 0, 0


    def next(self):
        """
        :rtype: int
        """
        while self.j >= len(self.v[self.i]):
            if self.i == 0: self.i = 1
            else: self.i, self.j = 0, self.j + 1
        ans = self.v[self.i][self.j]
        if self.i == 0: self.i += 1
        else: self.i, self.j = 0, self.j + 1
        return ans


    def hasNext(self):
        """
        :rtype: bool
        """
        if self.i == 1 and self.j >= len(self.v[0])-1 and self.j >= len(self.v[1]): return False
        if self.j >= len(self.v[0]) and self.j >= len(self.v[1]): return False
        return True