扁平化嵌套列表迭代器

题目描述

来源于 https://leetcode-cn.com/

给定一个嵌套的整型列表。设计一个迭代器,使其能够遍历这个整型列表中的所有整数。

列表中的项或者为一个整数,或者是另一个列表。

示例 1:

输入: [[1,1],2,[1,1]]
输出: [1,1,2,1,1]
解释: 通过重复调用 next 直到 hasNext 返回false,next 返回的元素的顺序应该是: [1,1,2,1,1]

示例 2:

输入: [1,[4,[6]]]
输出: [1,4,6]
解释: 通过重复调用 next 直到 hasNext 返回false,next 返回的元素的顺序应该是: [1,4,6]

解法:

class NestedIterator {
public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        push_nested_list_into_stack(nestedList);
        flatten_top();
    }
    
    int next() {
        int n = stk.top()->getInteger();
        stk.pop();
        return n;
    }
    
    bool hasNext() {
        flatten_top();
        return !stk.empty();
    }
private:
    void push_nested_list_into_stack(const vector<NestedInteger> & nestedList){
        for(int i=nestedList.size()-1;i>=0;i--){
            stk.push(&nestedList[i]);
        }
    }
    void flatten_top(){
        while(!stk.empty() && !stk.top()->isInteger()){
            const auto & nestedList = stk.top()->getList();
            stk.pop();
            push_nested_list_into_stack(nestedList);
        }
    }
    stack<const NestedInteger*> stk;
};