LeetCode 341: Flatten Nested List Iterator

link

We want to flatten the list lazily, so we use a stack with the top as the next element to return. If hasNext() returns true, the next() must return an integer. So, hasNext() must guarantee that there is an integer to return. We could not simply use len() as the signal for hasNext() because, a list is non-empty even if its only element is an empty list. So, in hasNext() we flatten the stack enough to make sure the top is an integer.

Time: \mathcal{O}(n), space: same.

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

class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        self.stack = [ x for x in reversed(nestedList) ]
    
    def next(self) -> int:
        return self.stack.pop().getInteger()
    
    def hasNext(self) -> bool:
        while self.stack and not self.stack[-1].isInteger():
            nested_list = self.stack.pop().getList()
            while nested_list:
                self.stack.append( nested_list.pop() )

        return len(self.stack) > 0

# Your NestedIterator object will be instantiated and called as such:
# i, v = NestedIterator(nestedList), []
# while i.hasNext(): v.append(i.next())

Leave a comment