Flatten Nested List Iterator
MedPrompt
You are given a nested list of integers nestedList. Each element is either an integer or a list whose elements may also be integers or other lists. Implement an iterator to flatten it.
Implement the NestedIterator class:
NestedIterator(List<NestedInteger> nestedList)Initializes the iterator with the nested listnestedList.int next()Returns the next integer in the nested list.boolean hasNext()Returnstrueif there are still some integers in the nested list andfalseotherwise.
Your code will be tested with the following pseudocode:
initialize iterator with nestedList
res = []
while iterator.hasNext()
append iterator.next() to the end of res
return resIf res matches the expected flattened list, then your code will be judged as correct.
Example 1:
Input: nestedList = [[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: nestedList = [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].
Constraints:
1 <= nestedList.length <= 500- The values of the integers in the nested list is in the range
[-106, 106].
Approaches
2 approaches with complexity analysis and trade-offs.
This approach involves pre-processing the entire nested list during the iterator's initialization. We traverse the nested structure recursively and store all the integers in a simple flat list. The hasNext() and next() methods then operate on this pre-computed list.
Algorithm
- Initialize an empty
ArrayList<Integer>calledflattenedListand an integercurrentIndexto 0. - Create a recursive helper method
flatten(List<NestedInteger> list). - Inside
flatten, iterate through eachNestedIntegerin the inputlist. - If the element
isInteger(), add its value toflattenedList. - If the element is a list, recursively call
flattenongetList(). - In the constructor, call
flattenwith the initialnestedList. - For
hasNext(), returntrueifcurrentIndexis less than the size offlattenedList,falseotherwise. - For
next(), return the element atcurrentIndexfromflattenedListand incrementcurrentIndex.
Walkthrough
In the constructor, we initialize an empty list of integers (e.g., ArrayList<Integer>). We then define a helper function, typically a recursive one, that takes a List<NestedInteger>. This helper function iterates through the given list. For each NestedInteger element, if it's an integer, we add it directly to our flat list. If it's a list, we make a recursive call to the helper function with this sublist. After the constructor finishes, we have a complete, flattened list of all integers. The hasNext() method simply checks if we have reached the end of our flat list, and the next() method retrieves the integer at the current position and advances the position.
/** * // This is the interface that allows for creating nested lists. * // You should not implement it, or speculate about its implementation * public interface NestedInteger { * * // @return true if this NestedInteger holds a single integer, rather than a nested list. * public boolean isInteger(); * * // @return the single integer that this NestedInteger holds, if it holds a single integer * // Return null if this NestedInteger holds a nested list * public Integer getInteger(); * * // @return the nested list that this NestedInteger holds, if it holds a nested list * // Return empty list if this NestedInteger holds a single integer * public List<NestedInteger> getList(); * } */public class NestedIterator implements Iterator<Integer> { private List<Integer> flattenedList; private int currentIndex; public NestedIterator(List<NestedInteger> nestedList) { this.flattenedList = new ArrayList<>(); this.currentIndex = 0; flatten(nestedList); } private void flatten(List<NestedInteger> list) { for (NestedInteger nestedInt : list) { if (nestedInt.isInteger()) { this.flattenedList.add(nestedInt.getInteger()); } else { flatten(nestedInt.getList()); } } } @Override public Integer next() { // The problem statement guarantees hasNext() will be called before next(). return this.flattenedList.get(currentIndex++); } @Override public boolean hasNext() { return this.currentIndex < this.flattenedList.size(); }}Complexity
Time
O(N + L) for the constructor, where N is the total number of integers and L is the total number of lists. We must visit every element to flatten the list. The `next()` and `hasNext()` methods are O(1) because they just involve an index check and an array access.
Space
O(N + D), where N is the total number of integers and D is the maximum nesting depth. O(N) space is required for storing the flattened list, and O(D) space is used by the recursion stack during the initial flattening process.
Trade-offs
Pros
Simple to implement and understand.
The
next()andhasNext()operations are extremely fast (O(1) worst-case time).
Cons
High memory usage as it requires storing all integers in memory at once. This can be a significant issue for very large nested lists.
There is a large upfront cost in the constructor to flatten the entire list before any iteration begins. This defeats the purpose of a lazy iterator.
May lead to a
StackOverflowErrorfor deeply nested lists due to recursion depth.
This approach avoids flattening the entire list upfront. Instead, it uses a stack to keep track of the traversal state. This is a form of iterative deepening search. We only process as much of the list as needed to find the next integer, making it a lazy iterator. This significantly improves space complexity, especially for large and sparse nested lists.
Algorithm
- Initialize a
Stack<NestedInteger>. - In the constructor, iterate through the input
nestedListfrom the last element to the first, pushing eachNestedIntegeronto the stack. - In
hasNext():- Start a loop that continues as long as the stack is not empty and the element at the top (
peek()) is a list. - Inside the loop,
pop()the list from the stack. - Iterate through this popped list from last to first, pushing each of its elements onto the stack.
- After the loop, return
!stack.isEmpty().
- Start a loop that continues as long as the stack is not empty and the element at the top (
- In
next():- First, ensure
hasNext()returns true (or assume it has been called). pop()the top element from the stack, get its integer value, and return it.
- First, ensure
Walkthrough
The core idea is to treat the nested list structure as a tree and perform a pre-order traversal iteratively. We use a Stack to store NestedInteger objects. In the constructor, we push the elements of the initial nestedList onto the stack in reverse order. This is crucial because a stack is a LIFO (Last-In, First-Out) structure, and we want to process the list elements from first to last. The hasNext() method is where the main logic resides. It ensures that the top of the stack is an integer. It repeatedly performs the following: while the stack is not empty and the top element is a list, pop the list. Then, push the elements of that popped list onto the stack, again in reverse order. This loop continues until the stack is empty or an integer is at the top. hasNext() then returns true if the stack is not empty, and false otherwise. The next() method assumes hasNext() has been called and an integer is ready. It simply pops the NestedInteger from the stack, retrieves its integer value, and returns it.
/** * // This is the interface that allows for creating nested lists. * // You should not implement it, or speculate about its implementation * public interface NestedInteger { * * // @return true if this NestedInteger holds a single integer, rather than a nested list. * public boolean isInteger(); * * // @return the single integer that this NestedInteger holds, if it holds a single integer * // Return null if this NestedInteger holds a nested list * public Integer getInteger(); * * // @return the nested list that this NestedInteger holds, if it holds a nested list * // Return empty list if this NestedInteger holds a single integer * public List<NestedInteger> getList(); * } */import java.util.NoSuchElementException; public class NestedIterator implements Iterator<Integer> { private Stack<NestedInteger> stack; public NestedIterator(List<NestedInteger> nestedList) { this.stack = new Stack<>(); // Push elements in reverse order to process them from start to end for (int i = nestedList.size() - 1; i >= 0; i--) { this.stack.push(nestedList.get(i)); } } @Override public Integer next() { if (!hasNext()) { throw new NoSuchElementException(); } // hasNext() ensures the top is an integer return stack.pop().getInteger(); } @Override public boolean hasNext() { // Ensure the top of the stack is an integer while (!stack.isEmpty() && !stack.peek().isInteger()) { List<NestedInteger> nestedList = stack.pop().getList(); for (int i = nestedList.size() - 1; i >= 0; i--) { stack.push(nestedList.get(i)); } } return !stack.isEmpty(); }}Complexity
Time
The cost is amortized over all calls. Each integer and each list is pushed onto and popped from the stack exactly once. Therefore, the total time complexity for iterating through all elements is O(N + L), where N is the total number of integers and L is the total number of lists. This means the amortized time complexity for each call to `hasNext()` and `next()` is O(1).
Space
O(D + M), where D is the maximum nesting depth and M is the maximum number of items in any single list. In the worst case, such as a list like `[[1], [2], [3], ...]`, the space can be O(L). However, for deeply nested lists like `[1, [2, [3, ...]]]` the space is O(D). This is generally much better than the O(N) space required by the first approach.
Trade-offs
Pros
It's a lazy iterator; computation is done on-demand, not all at once.
Memory efficient. The space used depends on the depth and breadth of the nesting, not the total number of integers, which is a significant advantage for large inputs.
Avoids the risk of
StackOverflowErrorby managing the stack explicitly instead of using system call stack (recursion).
Cons
The implementation is more complex than the pre-computation approach.
A single call to
hasNext()can be slow if it needs to unpack many nested lists to find the next integer. However, the cost is amortized.
Solutions
Solution
/** * // This is the interface that allows for creating nested lists. * // You should not implement it, or speculate about its implementation * public interface NestedInteger { * * // @return true if this NestedInteger holds a single integer, rather than a nested list. * public boolean isInteger(); * * // @return the single integer that this NestedInteger holds, if it holds a single integer * // Return null if this NestedInteger holds a nested list * public Integer getInteger(); * * // @return the nested list that this NestedInteger holds, if it holds a nested list * // Return null if this NestedInteger holds a single integer * public List<NestedInteger> getList(); * } */ public class NestedIterator implements Iterator < Integer > { private List < Integer > vals ; private Iterator < Integer > cur ; public NestedIterator ( List < NestedInteger > nestedList ) { vals = new ArrayList <>(); dfs ( nestedList ); cur = vals . iterator (); } @Override public Integer next () { return cur . next (); } @Override public boolean hasNext () { return cur . hasNext (); } private void dfs ( List < NestedInteger > nestedList ) { for ( NestedInteger e : nestedList ) { if ( e . isInteger ()) { vals . add ( e . getInteger ()); } else { dfs ( e . getList ()); } } } } /** * Your NestedIterator object will be instantiated and called as such: * NestedIterator i = new NestedIterator(nestedList); * while (i.hasNext()) v[f()] = i.next(); *//** * // This is the interface that allows for creating nested lists. * // You should not implement it, or speculate about its implementation * class NestedInteger { * public: * // Return true if this NestedInteger holds a single integer, rather than a nested list. * bool isInteger() const; * * // Return the single integer that this NestedInteger holds, if it holds a single integer * // The result is undefined if this NestedInteger holds a nested list * int getInteger() const; * * // Return the nested list that this NestedInteger holds, if it holds a nested list * // The result is undefined if this NestedInteger holds a single integer * const vector<NestedInteger> &getList() const; * }; */ class NestedIterator { public: NestedIterator ( vector < NestedInteger >& nestedList ) { dfs ( nestedList ); } int next () { return vals [ cur ++ ]; } bool hasNext () { return cur < vals . size (); } private: vector < int > vals ; int cur = 0 ; void dfs ( vector < NestedInteger >& nestedList ) { for ( auto & e : nestedList ) { if ( e . isInteger ()) { vals . push_back ( e . getInteger ()); } else { dfs ( e . getList ()); } } } }; /** * Your NestedIterator object will be instantiated and called as such: * NestedIterator i(nestedList); * while (i.hasNext()) cout << i.next(); */# """ # 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 # """ from collections import deque class NestedIterator : def __init__ ( self , nestedList : [ NestedInteger ]): self . stack = deque () self . prepareStack ( nestedList ) def next ( self ) -> int : if not self . hasNext (): # trigger hasNext() return None self . hasNext () return self . stack . pop (). getInteger () def hasNext ( self ) -> bool : while self . stack and not self . stack [ - 1 ]. isInteger (): # getList() more like get item, could be Integer or NestedInteger lst = self . stack . pop (). getList () self . prepareStack ( lst ) return bool ( self . stack ) def prepareStack ( self , nestedList ): for i in range ( len ( nestedList ) - 1 , - 1 , - 1 ): self . stack . append ( nestedList [ i ]) # Your NestedIterator object will be instantiated and called as such: # i, v = NestedIterator(nestedList), [] # while i.hasNext(): v.append(i.next()) ############ ''' >>> from collections import deque >>> >>> stack = deque() >>> stack.append(3) >>> stack.append(2) >>> stack.append(1) >>> stack deque([3, 2, 1]) >>> stack.pop() 1 ''' class NestedIterator : # not working if memory is limited and input is huge list def __init__ ( self , nestedList : [ NestedInteger ]): def dfs ( nestedList ): for e in nestedList : if e . isInteger (): self . vals . append ( e . getInteger ()) else : dfs ( e . getList ()) self . vals = [] dfs ( nestedList ) self . cur = 0 def next ( self ) -> int : res = self . vals [ self . cur ] self . cur += 1 return res def hasNext ( self ) -> bool : return self . cur < len ( self . vals ) # Your NestedIterator object will be instantiated and called as such: # i, v = NestedIterator(nestedList), [] # while i.hasNext(): v.append(i.next())Video walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.