Python, O(k) space, easy to extend to k case


  • -1
    B

    Basically using bookkeep data to keep track current progress.
    Once any row is exhausted, update the bookkeep data accordingly.

    class ZigzagIterator(object):

    def __init__(self, v1, v2):
        self.row = 0
        self.v = [v1, v2]
        self.k = len(self.v)
        self.q = [0] * self.k
            
    
    def next(self):
        #Update bookkeeping value, if one row is exhausted
        if self.q[self.row] == len(self.v[self.row]):
            self.v.pop(self.row)
            self.k -= 1
            self.row = (self.row + 1) % self.k
    
        #Output current, and prepare for next            
        x = self.v[self.row][self.q[self.row]]
        self.q[self.row] += 1
        self.row = (self.row + 1) % self.k
            
        return x
        
    
    def hasNext(self):
        return any([self.q[i] < len(self.v[i]) for i in xrange(self.k)])

Log in to reply
 

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