Simple python solution for 2 and multiple input (with comments)


  • 0
    S
    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.cLen = min(len(self.v1),len(self.v2))  # common length
            self.length = len(v1)+len(v2)  # total length
            self.ptr = 0        
    
        def next(self):
            """
            :rtype: int
            """
            if self.ptr<self.cLen*2:
                res = self.v1[self.ptr/2] if self.ptr%2==0 else self.v2[self.ptr/2]
            else:
                res = self.v1[self.ptr-self.cLen] if len(self.v1)>len(self.v2) else self.v2[self.ptr-self.cLen]
            self.ptr+=1
            return res
    
        def hasNext(self):
            """
            :rtype: bool
            """
            return self.ptr<self.length
    
    class ZigzagIterator2(object):
    
        def __init__(self, *vs):
            self.vs = list(vs)
            self.cLen = min(len(v) for v in self.vs)  # common length
            self.length = sum(len(v) for v in self.vs)  # total length
            self.ptr = 0
    
        def next(self):
            """
            :rtype: int
            """
            if self.ptr>=self.cLen*len(self.vs):  # if ptr reach the end of common length
                new_index = [i for i in xrange(len(self.vs)) if len(self.vs[i])>self.cLen]  # the index in the updated vs
                self.length -= self.cLen*(len(self.vs)-len(new_index))  # update the total length
                self.ptr -= self.cLen*(len(self.vs)-len(new_index))  # update the pos of ptr
                self.vs = [self.vs[i] for i in new_index]  # remove the shortest v
                self.cLen = min(len(v) for v in self.vs)
            index = self.ptr%len(self.vs)
            number = self.ptr/len(self.vs)
            res = self.vs[index][number]
            self.ptr+=1
            return res
    
        def hasNext(self):
            """
            :rtype: bool
            """
            return self.ptr<self.length
    
    # Your ZigzagIterator object will be instantiated and called as such:
    i, v = ZigzagIterator([1,2,3,4,5,6,7],[1,2,3,4,5]), []
    while i.hasNext(): v.append(i.next())
    print v
    i, v = ZigzagIterator2([1,2,3,4,5,6,7],[1,2,3,4,5],[1,2,3]), []
    while i.hasNext(): v.append(i.next())
    print v

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.