341. Flatten Nested List Iterator

Posted 2020-07-31 蜃利的阴影下

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:
Given the list [[1,1],2,[1,1]],

By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,1,2,1,1].

 

Example 2:
Given the list [1,[4,[6]]],

By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,4,6]

 

 

Time Limit Exceeded.

/**
 * // 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> {
    
    Deque<Integer> indices = new LinkedList<Integer>();
    Deque<List<NestedInteger>> levels = new LinkedList<>();
    
    /**
     * NestedInteger [1] is treated as an list.
     * [1].isInteger() == false;
     * 
     * NestedInteger [] is possible.
     */ 
    public NestedIterator(List<NestedInteger> nestedList) {
        while(nestedList.size() > 0)
        {
            indices.push(0);
            levels.push(nestedList);
            
            NestedInteger firstItem = nestedList.get(0);
            if(firstItem.isInteger())
                break;
            else
                nestedList = firstItem.getList();
        }
    }

    @Override
    public Integer next() {
        int currentInd = indices.pop();
        List<NestedInteger> currentLevel = levels.peek();
        Integer ret = currentLevel.get(currentInd).getInteger();
        
        if(currentInd+1<currentLevel.size())
            indices.push(currentInd+1);
        else
        {
            levels.pop(); //current level reaches the end. pop out current level.
            if(levels.isEmpty()) //if levels are empty, then there is no next.
                return ret;
            int indexForLastLevel = indices.pop();
            //do the loop to reach a level that is not at the end and not at an empty element.
            while(true)
            {
                currentLevel = levels.peek();
                while(indexForLastLevel+1<currentLevel.size())
                {
                    NestedInteger currentItem = currentLevel.get(indexForLastLevel+1);
                    if(!currentItem.isInteger() && currentItem.getList().isEmpty())
                        ++indexForLastLevel;  //skip the empty list items.
                }
                //reach the end of current level. Go the upper level.
                if(indexForLastLevel+1 == currentLevel.size())
                {
                    levels.pop();
                    if(levels.isEmpty())  //reached the root.
                        return ret;
                    indexForLastLevel = indices.pop();
                }
                else
                    break;
            }
            
            indices.push(indexForLastLevel+1);
            NestedInteger currentItem = levels.peek().get(indexForLastLevel+1);
            while(!currentItem.isInteger())
            {
                indices.push(0);
                List<NestedInteger> nextLevel = currentItem.getList();
                levels.push(nextLevel);
                currentItem = nextLevel.get(0);
            }
            
        }
        return ret;
    }

    @Override
    public boolean hasNext() {
        if(indices.size() > 0)
            return true;
        return false;
    }
}

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i = new NestedIterator(nestedList);
 * while (i.hasNext()) v[f()] = i.next();
 */