Problem statement

https://leetcode.com/problems/zigzag-iterator/

Solution

Define self.v = [v1, v2], self.p are indexes in these arrays and self.turn is turn which array we choose. Our invariant that we are always on the place and turn that we will choose.

  1. To check hasNext we just check if we not at the end of both arrays.
  2. To find next we check if we can take element from given point. If we can, move this pointer and change turn. If we can not, change turn and then move pointer.

Complexity

Time complexity is just $O(1)$ for each operation. Space complexity for whole data structure is $O(l_1 + l_2)$, where $l_1$ and $l_2$ are lengths of v1 and v2.

Code

class ZigzagIterator:
    def __init__(self, v1, v2):
        self.v = [v1, v2]
        self.p, self.turn = [0, 0], 0
        
    def next(self):
        if self.p[self.turn] < len(self.v[self.turn]):
            ans = self.v[self.turn][self.p[self.turn]]
            self.p[self.turn] += 1
            self.turn = (self.turn + 1)%2
        else:
            self.turn = (self.turn + 1)%2
            ans = self.v[self.turn][self.p[self.turn]]
            self.p[self.turn] += 1
            
        return ans    

    def hasNext(self):
        return self.p[0] < len(self.v[0]) or self.p[1] < len(self.v[1])