LeetCode: 341. Flatten Nested List Iterator

Making a new list is an easy solution to implement with recursion.

class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        self.res = []
        self._flattener(nestedList)
        self.ptr = 0
            
    def _flattener(self, nestedList):
        for nl in nestedList:
            if nl.isInteger():
                self.res.append(nl.getInteger())
            else:
                self._flattener(nl.getList())
    
    def next(self) -> int:
        ret = self.res[self.ptr]
        self.ptr += 1
        return ret
    
    def hasNext(self) -> bool:
        return self.ptr < len(self.res)
         

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

 

However, above does not seem to act like an real iterator, as it is not lazy.

We can use 'yield' to make something really works like an iterator as below.

class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        def _generator(nestedList):
            for nl in nestedList:
                if nl.isInteger():
                    yield nl.getInteger()
                else:
                    yield from _generator(nl.getList())
        self.gen = _generator(nestedList)
        self.nextelem = next(self.gen, None)
    
    def next(self) -> int:
        ret = self.nextelem
        self.nextelem = next(self.gen, None)
        return ret
    
    def hasNext(self) -> bool:
        return self.nextelem != None