Problem

You are given a nested list of integer,s 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 list nestedList.
  • int next() Returns the next integer in the nested list.
  • boolean hasNext() Returns true if there are still some integers in the nested list and false otherwise.

Your code will be tested with the following pseudocode:

1
2
3
4
5

initialize iterator with nestedList
res = []
while iterator.hasNext()
    append iterator.next() to the end of res
return res

If res matches the expected flattened list, then your code will be judged as correct.

Examples

Example 1:

1
2
3

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:

1
2
3

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].

Solution

To flatten a nested list, you can either use recursion to flatten it completely and then iterate over a flat list or use a stack to manage the nested structure iteratively.

Method 1 - Recursion

Above approach is good, but it is very eager in nature, where we are eagerly creating a list of flattened integers and then returning. Here is the approach:

  • Flatten the nested list recursively.
  • Store the flattened list in a simple list structure.
  • Iterate over the flattened list using index.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

public class NestedIterator implements Iterator<Integer> {
    private List<Integer> flatList;
    private int curr, size;

    public NestedIterator(List<NestedInteger> nestedList) {
        this.flatList = new ArrayList<>();
        curr = 0;
        flatten(nestedList);
        size = flatList.size();
    }

    @Override
    public Integer next() {
        return curr < size ? flatList.get(curr++) : -1;
    }

    @Override
    public boolean hasNext() {
        return curr < size;
    }
    
    private void flatten(List<NestedInteger> nestedList) {
        for (NestedInteger n : nestedList) {
            if (n.isInteger()) {
                flatList.add(n.getInteger());
            } else {
                flatten(n.getList());
            }
        }
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

class NestedInteger:
    def isInteger(self) -> bool:
        ...

    def getInteger(self) -> int:
        ...

    def getList(self) -> List['NestedInteger']:
        ...

class NestedIterator:
    def __init__(self, nestedList: List[NestedInteger]) -> None:
        self.flat_list: List[int] = []
        self.index: int = 0
        self.flatten(nestedList)

    def flatten(self, nestedList: List[NestedInteger]) -> None:
        for ni in nestedList:
            if ni.isInteger():
                self.flat_list.append(ni.getInteger())
            else:
                self.flatten(ni.getList())

    def next(self) -> int:
        ans = self.flat_list[self.index]
        self.index += 1
        return ans

    def hasNext(self) -> bool:
        return self.index < len(self.flat_list)

Method 2 - Using Stack

In the constructor, elements from nestedList are inserted into the stack in reverse order to ensure the first element can be accessed first when popped. In the hasNext() method, the current top of the stack is checked: if it is an integer, it is returned and removed from the stack; if it is a list, it is further flattened. This method iteratively flattens the nested list by processing elements from back to front.

Here is the approach:

  • Use a stack to keep track of elements.
  • Push elements onto the stack from back to front.
  • Flatten nested lists iteratively by checking the top of the stack and processing lists as needed.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

public class NestedIterator implements Iterator<Integer> {
    Deque<NestedInteger> stack = new ArrayDeque<>();

    public NestedIterator(List<NestedInteger> nestedList) {
        flatten(nestedList);
    }

    @Override
    public Integer next() {
        if (!hasNext()) {
            return null;
        }
        return stack.pop().getInteger();
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty() && !stack.peek().isInteger()) {
            List<NestedInteger> list = stack.pop().getList();
            flatten(list);
        }
        return !stack.isEmpty();
    }
    
    private void flatten(List<NestedInteger> nestedList) {
        for (int i = nestedList.size() - 1; i >= 0; i--) {
            stack.push(nestedList.get(i));
        }
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

class NestedInteger:
    def isInteger(self) -> bool:
        ...

    def getInteger(self) -> int:
        ...

    def getList(self) -> List['NestedInteger']:
        ...

class NestedIterator:
    def __init__(self, nestedList: List[NestedInteger]) -> None:
        self.stack: List[NestedInteger] = []
        self.flatten(nestedList)

    def flatten(self, nestedList: List[NestedInteger]) -> None:
        for i in range(len(nestedList) - 1, -1, -1):
            self.stack.append(nestedList[i])

    def next(self) -> int:
        return self.stack.pop().getInteger()

    def hasNext(self) -> bool:
        while self.stack and not self.stack[-1].isInteger():
            lst = self.stack.pop().getList()
            self.flatten(lst)
        return bool(self.stack)

Complexity

  • ⏰ Time complexity: O(n), where n is the total number of integers in all nested lists, since every element is processed once.
  • 🧺 Space complexity: O(NNNXXX)