341. Flatten Nested List Iterator

Problem

Given a nested list of integers, implement an iterator to flatten it.

Each element is either an integer, or a list -- whose elements may also be integers or other lists.

Example 1:

Input:[[1,1],2,[1,1]]
Output:[1,1,2,1,1]
Explanation:By callingnext repeatedly untilhasNext returns false,
             the order of elements returned bynext should be:[1,1,2,1,1].

Example 2:

Input:[1,[4,[6]]]
Output:[1,4,6]
Explanation:By callingnext repeatedly untilhasNext returns false,
             the order of elements returned bynext should be:[1,4,6].

Problem Summary

Given a nested list of integers, design an iterator so that it can traverse all integers in this integer list. Each item in the list is either an integer or another list. The elements of the list may also be integers or other lists.

Solution Approach

The problem requires implementing a nested version of an array. It can be implemented with []int, or with a linked list. Here, the author implements it with a linked list. Construct a one-dimensional array on the outside, where each element inside the one-dimensional array is a linked list. Additionally, record the index of this nested linked list in the original array. The implementation of Next() is relatively simple: retrieve the corresponding nested node. The HasNext() method determines whether there is a next element based on the index information inside the nested node.

Code

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * type NestedInteger struct {
 * }
 *
 * // Return true if this NestedInteger holds a single integer, rather than a nested list.
 * func (this NestedInteger) IsInteger() bool {}
 *
 * // 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
 * // So before calling this method, you should have a check
 * func (this NestedInteger) GetInteger() int {}
 *
 * // Set this NestedInteger to hold a single integer.
 * func (n *NestedInteger) SetInteger(value int) {}
 *
 * // Set this NestedInteger to hold a nested list and adds a nested integer to it.
 * func (this *NestedInteger) Add(elem NestedInteger) {}
 *
 * // Return the nested list that this NestedInteger holds, if it holds a nested list
 * // The list length is zero if this NestedInteger holds a single integer
 * // You can access NestedInteger's List element directly if you want to modify it
 * func (this NestedInteger) GetList() []*NestedInteger {}
 */

type NestedIterator struct {
    stack *list.List
}

type listIndex struct {
    nestedList []*NestedInteger
    index int
}

func Constructor(nestedList []*NestedInteger) *NestedIterator {
    stack := list.New()
    stack.PushBack(&listIndex{nestedList, 0})
    return &NestedIterator{stack}
}

func (this *NestedIterator) Next() int {
    if !this.HasNext() {
        return -1
    }
    last := this.stack.Back().Value.(*listIndex)
    nestedList, i := last.nestedList, last.index
    val := nestedList[i].GetInteger()
    last.index++
    return val
}

func (this *NestedIterator) HasNext() bool {
    stack := this.stack
    for stack.Len() > 0 {
        last := stack.Back().Value.(*listIndex)
        nestedList, i := last.nestedList, last.index
        if i >= len(nestedList) {
            stack.Remove(stack.Back())
        } else {
            val := nestedList[i]
            if val.IsInteger() {
                return true
            }
            last.index++
            stack.PushBack(&listIndex{val.GetList(), 0})
        }
    }
    return false
}