leetcode 341. Flatten Nested List Iterator

Posted 2020-07-18 godjob

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].

Subscribe to see which companies asked this question

 

很显然这个是一个递归问题，但是题目要求实现一个迭代器，这里我想到两个方法：

1. 直接在构造函数中使用递归构造好整个序列，然后逐个返回

2.栈保存历史记录，栈顶保存当前打印的列表以及打印到列表哪个位置，刚开始实现很顺利，但是后面提交发现测试用例有空列表的情况，如输入为[[],[[]]]，这个

就麻烦了，该进后的方式就是保存一个缓存节点

下面是自己实现的方法：

 

#include <stack>
#include <vector>
#include <assert.h>
struct NestedContext {
    const vector<NestedInteger>* vec;
    size_t pos;
};

class NestedIterator {
public:
    NestedIterator(vector<NestedInteger> &nestedList)
    {
        has_next = true;
        if (nestedList.empty()) {
            return;
        }
        NestedContext ctx = {&nestedList, 0};
        s.push(ctx);
    }
    int next() {
        return cache_val;
    }

    bool _next(int& val) {
        while (!s.empty()) {
            NestedContext& ctx = s.top();
            if (ctx.vec->at(ctx.pos).isInteger()) {
                int rv = ctx.vec->at(ctx.pos++).getInteger();
                if (ctx.pos >= ctx.vec->size()) {
                    s.pop();
                }
                val = rv;
                return true;
            }
            else {
                const vector<NestedInteger>& nestedList = ctx.vec->at(ctx.pos++).getList();
                if (ctx.pos >= ctx.vec->size()) {
                    s.pop();
                }
                if (!nestedList.empty()) {
                    NestedContext new_ctx = {&nestedList, 0};
                    s.push(new_ctx);
                }
            }
        }
        return false;
    }

    bool hasNext() {
        has_next = _next(cache_val);
        return !s.empty() || has_next;
    }
    stack<NestedContext> s;
    bool has_next;
    int cache_val;
};