341. Flatten Nested List Iterator

341. Flatten Nested List Iterator (Medium)

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 calling next repeatedly until hasNext returns false, 
             the order of elements returned by next should be: [1,1,2,1,1].

Example 2:

Input: [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].

Related Topics:
Stack, Design

Similar Questions:

• Flatten 2D Vector (Medium)
• Zigzag Iterator (Medium)
• Mini Parser (Medium)
• Array Nesting (Medium)

Solution 1.

// OJ: https://leetcode.com/problems/flatten-nested-list-iterator/
// Author: github.com/lzl124631x
// Time: O(1) amortized
// Space: O(D) where D is the max depth of the list
class NestedIterator {
    typedef vector<NestedInteger>::iterator iter;
    stack<pair<iter, iter>> s;
    void goToInteger() {
        while (s.size()) {
            if (s.top().first == s.top().second) {
                s.pop();
                if (s.size()) s.top().first++;
            } else if (s.top().first->isInteger()) break;
            else {
                auto &list = s.top().first->getList();
                s.emplace(list.begin(), list.end());
            }
        }
    }
public:
    NestedIterator(vector<NestedInteger> &list) {
        s.emplace(list.begin(), list.end());
        goToInteger();
    }
    
    int next() {
        int val = s.top().first->getInteger();
        s.top().first++;
        goToInteger();
        return val;
    }
    
    bool hasNext() {
        return s.size();
    }
};

Solution 2.

// OJ: https://leetcode.com/problems/flatten-nested-list-iterator/
// Author: github.com/lzl124631x
// Time: O(1) amortized
// Space: O(N)
class NestedIterator {
    stack<NestedInteger> s;
    void goToInteger() {
        while (s.size() && !s.top().isInteger()) {
            auto list = s.top().getList();
            s.pop();
            for (int i = list.size() - 1; i >= 0; --i) s.push(list[i]);
        }
    }
public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        for (int i = nestedList.size() - 1; i >= 0; --i) s.push(nestedList[i]);
        goToInteger();
    }
    int next() {
        int val = s.top().getInteger();
        s.pop();
        goToInteger();
        return val;
    }
    bool hasNext() {
        return s.size();
    }
};