Flatten 2D Vector

Problem: Flatten 2D Vector

The problem can be solved by two methods.

First, you can use traditional way to use x and y to point to current node. If you do this way, you better add a method to make the pointer valid, which can be both called during initiating and calling of hasNext() method.

Code in Python:

class Vector2D(object):

    def __init__(self, vec2d):
        """
        Initialize your data structure here.
        :type vec2d: List[List[int]]
        """
        self.vec2d = vec2d
        self.x = self.y = 0
        self.makeValid()


    def makeValid(self):
        while self.x >= len(self.vec2d) or self.y >= len(self.vec2d[self.x]):
            if self.x >= len(self.vec2d): return False
            else: self.x, self.y = self.x+1, 0
        return True


    def next(self):
        """
        :rtype: int
        """
        ans = self.vec2d[self.x][self.y]
        if self.y < len(self.vec2d[self.x])-1: self.y += 1
        else: self.x, self.y = self.x+1, 0
        return ans


    def hasNext(self):
        """
        :rtype: bool
        """
        return self.makeValid()

Another method is just to store the flattened vector.

Code in Python:

class Vector2D(object):

    def __init__(self, vec2d):
        """
        Initialize your data structure here.
        :type vec2d: List[List[int]]
        """
        self.vec = [x for vec in vec2d for x in vec]
        self.index = 0


    def next(self):
        """
        :rtype: int
        """
        self.index += 1
        return self.vec[self.index-1]


    def hasNext(self):
        """
        :rtype: bool
        """
        return self.index < len(self.vec)