251. Flatten 2D Vector

JIAKAOBO

LeetCode

感谢赞助!

• ㊗️
• 大家
• offer
• 多多！

Problem

Design and implement an iterator to flatten a 2d vector. It should support the following operations: next and hasNext.

Example:

Vector2D iterator = new Vector2D([[1,2],[3],[4]]);

iterator.next(); // return 1
iterator.next(); // return 2
iterator.next(); // return 3
iterator.hasNext(); // return true
iterator.hasNext(); // return true
iterator.next(); // return 4
iterator.hasNext(); // return false

Notes:

Please remember to RESET your class variables declared in Vector2D, as static/class variables are persisted across multiple test cases. Please see here for more details.
You may assume that next() call will always be valid, that is, there will be at least a next element in the 2d vector when next() is called.

Code

class Vector2D {
    int index1 = 0;
    int index2 = 0;
    int[][] nums;
    public Vector2D(int[][] v) {
        nums = v;
    }

    public int next() {
        hasNext();
        return nums[index1][index2++];
    }

    public boolean hasNext() {
        while(index1 < nums.length) {
            if(index2 < nums[index1].length){
                return true;
            }

            index1++;
            index2 = 0;
        }

        return false;
    }
}

class Vector2D:

    def __init__(self, v: List[List[int]]):
        self.vector = v
        self.inner = 0
        self.outer = 0

    def advance_to_next(self):
        while self.outer < len(self.vector) and self.inner == len(self.vector[self.outer]):
            self.outer += 1
            self.inner = 0

    def next(self) -> int:
        self.advance_to_next()
        result = self.vector[self.outer][self.inner]
        self.inner += 1
        return result


    def hasNext(self) -> bool:
        self.advance_to_next()
        return self.outer < len(self.vector)