diff options
Diffstat (limited to 'flatten_nested_iterator.py')
-rw-r--r-- | flatten_nested_iterator.py | 54 |
1 files changed, 54 insertions, 0 deletions
diff --git a/flatten_nested_iterator.py b/flatten_nested_iterator.py new file mode 100644 index 0000000..1c9042c --- /dev/null +++ b/flatten_nested_iterator.py @@ -0,0 +1,54 @@ +# """ +# 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): + self.stack = [[nestedList, 0]] + + def next(self): + self.hasNext() + nestedList, i = self.stack[-1] + self.stack[-1][1] += 1 + return nestedList[i].getInteger() + + def hasNext(self): + s = self.stack + while s: + nestedList, i = s[-1] + if i == len(nestedList): + s.pop() + else: + x = nestedList[i] + if x.isInteger(): + return True + s[-1][1] += 1 + s.append([x.getList(), 0]) + return False + + +# Your NestedIterator object will be instantiated and called as such: +# i, v = NestedIterator(nestedList), [] +# while i.hasNext(): v.append(i.next()) |