341. Flatten Nested List Iterator

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].
class NestedIterator {
public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        for (int i = nestedList.size() - 1; i >= 0; --i) {
            st.push(nestedList[i]);
        }
    }

    int next() {
        NestedInteger t = st.top(); st.pop();
        return t.getInteger();
    }

    bool hasNext() {
        while (!st.empty()) {
            NestedInteger t = st.top();
            if (t.isInteger()) return true;
            st.pop();
            vector<NestedInteger> list = t.getList();
            for (int i = list.size() - 1; i >= 0; --i) {
                st.push(list[i]);
            }
        }
        return false;
    }
    stack<NestedInteger> st;
};

Last updated

Was this helpful?