Flatten Nested List Iterator (Leetcode 341)

Leetcode 341

Difficulty : Medium Link : April Leetcoding Day 13

Problem Description

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) Initialises 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.


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


  • 1 <= nestedList.length <= 500
  • The values of the integers in the nested list is in the range [-106, 106].


The problem is very Swifty in nature because it you look carefully this is the exact what InteratorProtocol does. If you use a forin loop with an array, set, or any other collection or sequence, you’re already using that type’s iterator (with syntax sugar ofcourse).

let nums:[Int] = [1,2,3,4,5]
for num in nums {
// iterating throug the sequence

Behind the hood, the system creates an iterator using makeIterator() and repeatedly calls next() on the sequence . Since Sequence protocol have an associated type of IteratorProtocol, you never have to worry would how next element would be computed.

The first instinct might be use flatmap() operator directly but you should think about it again. All higher-order operators are defined on Sequence protocol which itself confirm to IteratorProtocol. Here it is the precondition that you have to sort of implement the Iterator functionality itself. So if you are given this problem in an interview, be sure to check with your interviewer first.

The solution is simple, we can have a property for storing list of integers and a index property to maintain the current index. Upon initialisation, we can recursively traverse the nested lists in a depth-first manner. The next thing remains is to simply check if an element is an integer or a nested list. You might be temped to use conditional cast as? Int, however it wouldn’t work. Why? Because NestedInteger is a private class while Int is a struct. You can find that there are utility methods isInteger() and getInteger() which solves our problem. We can create a recursive function to append all integers and recursively calling itself on nested lists.

func flattenList(_ nestedList: [NestedInteger]) {
for element in nestedList {
if element.isInteger() {
view raw flatten_list.swift hosted with ❤ by GitHub

For getting the next element, we can simple use defer to increment the count after we have returned the current element. HasNext() is pretty straightforward as we only need to compare currentIndex with list count.

func next() -> Int {
defer {
currentIndex += 1
return list[currentIndex]



class NestedIterator {
private var list:[Int]
private var currentIndex: Int
init(_ nestedList: [NestedInteger]) {
self.list = []
currentIndex = 0
func flatten(_ nestedList: [NestedInteger]) {
for element in nestedList {
if element.isInteger() {
func next() -> Int {
defer {
currentIndex += 1
return list[currentIndex]
func hasNext() -> Bool {
return currentIndex < list.count


class NestedIterator:
def __init__(self, nestedList: [NestedInteger]):
def flatten(a):
req = []
for ele in a:
if ele.isInteger():
return req
self.flatten = flatten(nestedList)
self.curr = 0
def next(self) -> int:
t = self.flatten[self.curr]
return t
def hasNext(self) -> bool:
return self.curr<len(self.flatten)


public class NestedIterator implements Iterator<Integer> {
Queue<Integer> data = new LinkedList<>();
public NestedIterator(List<NestedInteger> nestedList) {
public void flatten(List<NestedInteger> list) {
for (NestedInteger el : list)
if (el.isInteger()) data.add(el.getInteger());
else flatten(el.getList());
public Integer next() { return data.poll(); }
public boolean hasNext() { return data.size() > 0; }


Time = O(N) : Where N is the total number of element

Space = O(N)

Leave a Reply

Your email address will not be published. Required fields are marked *