341. Flatten Nested List Iterator (Python)

Related Topic

Stack.

Description

Given a nested list of integers, implement an iterator to flatten it.

Each element is either an integer, or a list – whose elements may also be integers or other lists.

Sample I/O

Example 1

Input: [[1,1],2,[1,1]]
Output: [1,1,2,1,1]
Explanation: By calling next repeatedly until hasNext returns false, 
the order of elements returned by next should be: [1,1,2,1,1].

Example 2

Input: [1,[4,[6]]]
Output: [1,4,6]
Explanation: By calling next repeatedly until hasNext returns false, 
the order of elements returned by next should be: [1,4,6].

Methodology

The question offered interfaces to help us. We can use stack to store the elements in nestedList in reversed order. In the hasNext function we will pop out the top element of the stack. If the top element is integer, we will return True and let next function return the top integer. If the top element is not integer, which means it is the nestedList, we will use getList interface function to get the nested list that this nestedInteger holds, then concatenate it to the previous stack and assign it to stack, this will recusivelt get the integer. The next function will return the top integer of the stack.

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
#class NestedInteger(object):
#    def isInteger(self):
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        :rtype bool
#        """
#
#    def getInteger(self):
#        """
#        @return the single integer that this NestedInteger holds, if it holds a single integer
#        Return None if this NestedInteger holds a nested list
#        :rtype int
#        """
#
#    def getList(self):
#        """
#        @return the nested list that this NestedInteger holds, if it holds a nested list
#        Return None if this NestedInteger holds a single integer
#        :rtype List[NestedInteger]
#        """

class NestedIterator(object):

    def __init__(self, nestedList):
        """
        Initialize your data structure here.
        :type nestedList: List[NestedInteger]
        """
        self.stack = []
        while nestedList:
            self.stack.append(nestedList.pop())

    def next(self):
        """
        :rtype: int
        """
        return self.stack.pop().getInteger()     

    def hasNext(self):
        """
        :rtype: bool
        """
        while self.stack:
            top = self.stack[-1]
            if top.isInteger():
                return True
            self.stack = self.stack[:-1] + top.getList()[::-1]
        return False

BigO

Iterate the stack will take O(n+l) where n is the total number of integers and l is total number of lists