Skip to main content

341. Flatten Nested List Iterator

https://leetcode.com/problems/flatten-nested-list-iterator

Python

class NestedIterator:
def __init__(self, nestedList: [NestedInteger]):
self.cache = [item for item in self._flatten(nestedList)]
self.cur = 0
self.length = len(self.cache)

def _flatten(self, node):
if not node:
return

if isinstance(node, list):
items = node
elif node.isInteger():
items = []
yield node
else:
items = node.getList()

for item in items:
if item.isInteger():
yield item.getInteger()
else:
for children in item.getList():
for child in self._flatten(children):
yield child

def next(self) -> int:
i = self.cur
self.cur += 1
return self.cache[i]

def hasNext(self) -> bool:
return self.cur < self.length
# """
# 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
# """


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