341.flatten-nested-list-iterator.md

341. Flatten Nested List Iterator

NestedInteger is an interface that represents either an integer or a nested list. And the input is a list of NestedIntegers, it might be

// Input
[0, [1, 2], 3, [4, [5, 6], 7], 8]

// Its representation
[0, List  , 3, List          , 8] 
    [1, 2] 
               [4, List  , 7]
                   [5, 6]

// Its flattened representation
[0, 1, 2, 3, 4, 5, 6, 7, 8]

NOTE: [[]] is valid test case, it's nested list but has no integer, it should return false for hasNext().

Stack

We flatten the nested list on the fly, and use a stack to handle the nested case. We put the nested list in reverse order into the stack at the beginning, and flatten the list when we need to get the next integer.

Why putting the nested list in reverse order? Because we use stack to handle the nested case, and we want to pop the first element first, so that we can get the integer in the nested list first.

Then we can define the following functions:

• next(): Int: It return the integer, so we should have already prepared the first integer in the stack, that is, we should flatten the nested list and ensure the first element in the stack is integer.

For this function, it pops the first element from the stack, it should be integer, not the nested list. We don't flatten here because we might get the empty list after flatten in this function, and that will cause problem: hasNext() returns true but it's empty.

Stack = [0, ...]
         * // Integer, just return

• hasNext(): Boolean, we try flatten the nested list first, and push them to the stack, so that every item in the stack is an integer. If the stack is empty, then we know there is no more integer.

        [[1, 2], ...]
Stack = [List  , ...]
         * // List, flatten it and push to stack

// The stack after flatten
Stack = [1, 2  , ...]

// ----------------------------
// or the nested list is empty
        [[], ...]
Stack = [List  , ...]

// The stack after flatten
Stack = [] // Empty

class NestedIterator(private val nestedList: List<NestedInteger>) {

    private val stack = Stack<NestedInteger>()

    init {
        pushListIntoStack(nestedList)
    }

    fun next(): Int {
        return stack.pop().getInteger()!!
    }
    
    fun hasNext(): Boolean {
        // Flatten the nested item
        while (stack.isNotEmpty() && !stack.peek().isInteger()) {
            pushListIntoStack(stack.pop().getList()!!)
        }
        // And check if has next
        return stack.isNotEmpty()
    }

    // Flatten list and push to stack backwards
    private fun pushListIntoStack(list: List<NestedInteger>) {
        for (i in list.size - 1 downTo 0) {
            stack.push(list[i])
        }
    }
}

Queue

We pre-process every item at the beginning, we flatten all items and push into queue. We start by iterating every item, then:

• If it's integer, just enque.
• If it's list, flatten it by calling function recursively (implicit stack).

[0, [1, [2, 3], ...]
 * // Integer, just enque
     * // List, call flatten([1, [2, 3]) recursively
        * // Integer, enque
            * // List, call flatten([2, 3]) recursively
            * // Integer, enque
            * // Integer, enque

However, this approach is not recommended because it's not lazy, we flatten the whole list at the beginning.

class NestedIterator(private val nestedList: List<NestedInteger>) {

    private val queue = ArrayDeque<Int>()

    init {
        flatten(nestedList)
    }

    fun next(): Int {
        return queue.removeFirst()
    }
    
    fun hasNext(): Boolean {
        return queue.isNotEmpty()
    }

    private fun flatten(nestedList: List<NestedInteger>) {
        for (integer in nestedList) {
            if (integer.isInteger()) {
                queue.addLast(integer.getInteger()!!)
            } else {
                flatten(integer.getList()!!)
            }
        }
    }
}

Failed Case

My original implementation is to use a stack to handle nested case, and a queue to iterate the next integer at the current level.

[1, [2, 3], 4]

Queue = [1, List, 4]
         * // Integer, just return

• If we meet an integer, just return in next() function.
• If we meet the list, then we push the current queue to the stack, and set the queue to the new list of the nested list:

Stack = [4]
Queue = [2, 3]

But it failed on the empty list case [[]]. The problem is that we implement hasNext() wrong and we keep accessing the queue / stack even if they is empty.

class NestedIterator(private val nestedList: List<NestedInteger>) {

    private val stack = Stack<ArrayDeque<NestedInteger>>()
    private var queue = ArrayDeque<NestedInteger>(nestedList)

    fun next(): Int {
        if (queue.isNotEmpty()) {
            val nestedInteger = queue.removeFirst()
            if (nestedInteger.isInteger()) {
                return nestedInteger.getInteger()!!
            } else {
                stack.push(queue)
                queue = ArrayDeque(nestedInteger.getList()!!)
                return next()
            }
        } else {
            // The wrong line, the queue and stack might be empty here.
            queue = stack.pop()
            return next()
        }
    }
    
    fun hasNext(): Boolean {
        return queue.isNotEmpty() || (stack.isNotEmpty() && stack.peek().isNotEmpty())
    }
}