Leetcode: Zigzag Iterator

Question

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?
Hide Tags Design
Hide Similar Problems (M) Binary Search Tree Iterator (M) Flatten 2D Vector

My Solution

class ZigzagIterator(object):

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

        self.v1, self.v2 = v1, v2
        self.ind1, self.ind2 = 0,0

        self.flag = 1


    def next(self):
        """ :rtype: int """

        if self.hasNext():
            if self.ind1==len(self.v1) and self.ind2!=len(self.v2):
                res = self.v2[self.ind2]
                self.ind2 += 1
            elif self.ind2==len(self.v2) and self.ind1!=len(self.v1):
                res = self.v1[self.ind1]
                self.ind1 += 1
            else:
                if self.flag==1:
                    res = self.v1[self.ind1]
                    self.ind1 += 1
                    self.flag = 2
                else:
                    res = self.v2[self.ind2]
                    self.ind2 += 1
                    self.flag = 1
            return res                    

    def hasNext(self):
        """ :rtype: bool """

        if self.ind1==len(self.v1) and self.ind2==len(self.v2):
            return False
        else:
            return True



# Your ZigzagIterator object will be instantiated and called as such:
# i, v = ZigzagIterator(v1, v2), []
# while i.hasNext(): v.append(i.next())