LeetCode 341. 扁平化嵌套列表迭代器

题目描述

341. 扁平化嵌套列表迭代器

解法：递归

这道题最重要的是先理解题意，在理解的情况下，其实我们要做的就是一棵 N 叉树的遍历，为什么这样说呢？

比如说输入是 $[[1,1],2,[1,1]]$，其实就是如下树状结构：

把一个 NestedInteger 扁平化就等价于遍历一棵 N 叉树的所有「叶子节点」，我把所有叶子节点都拿出来，不就可以作为迭代器进行遍历了吗？

有了如此思路，下面的代码就容易解决了

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * class NestedInteger {
 *   public:
 *     // Return true if this NestedInteger holds a single integer, rather than a nested list.
 *     bool isInteger() const;
 *
 *     // Return the single integer that this NestedInteger holds, if it holds a single integer
 *     // The result is undefined if this NestedInteger holds a nested list
 *     int getInteger() const;
 *
 *     // Return the nested list that this NestedInteger holds, if it holds a nested list
 *     // The result is undefined if this NestedInteger holds a single integer
 *     const vector<NestedInteger> &getList() const;
 * };
 */

class NestedIterator {
private:
    vector<int>::iterator cur;
    vector<int> res;

    void traverse(NestedInteger root){
        if (root.isInteger())
        {
            res.push_back(root.getInteger());
            return;
        }
        for (NestedInteger child: root.getList())
            traverse(child);
    }

public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        for (NestedInteger node: nestedList)
            traverse(node);
        cur = res.begin();
    }
    
    int next() {
        return *cur++;
    }
    
    bool hasNext() {
        return cur != res.end();
    }
};

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i(nestedList);
 * while (i.hasNext()) cout << i.next();
 */

以上解法虽然可以通过，但是在面试中，也许是有瑕疵的

我们的解法中，一次性算出了所有叶子节点的值，全部装到 res 列表，也就是内存中，next 和 hasNext 方法只是在对 res 列表做迭代。如果输入的规模非常大，构造函数中的计算就会很慢，而且很占用内存

一般的迭代器求值应该是「惰性的」，也就是说，如果你要一个结果，我就算一个（或是一小部分）结果出来，而不是一次把所有结果都算出来

如果想做到这一点，使用递归函数进行 DFS 遍历肯定是不行的，而且我们其实只关心「叶子节点」，所以传统的 BFS 算法也不行。实际的思路很简单：

调用 hasNext 时，如果 nestedList 的第一个元素是列表类型，则不断展开这个元素，直到第一个元素是整数类型

由于调用 next 方法之前一定会调用 hasNext 方法，这就可以保证每次调用 next 方法的时候第一个元素是整数型，直接返回并删除第一个元素即可

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * class NestedInteger {
 *   public:
 *     // Return true if this NestedInteger holds a single integer, rather than a nested list.
 *     bool isInteger() const;
 *
 *     // Return the single integer that this NestedInteger holds, if it holds a single integer
 *     // The result is undefined if this NestedInteger holds a nested list
 *     int getInteger() const;
 *
 *     // Return the nested list that this NestedInteger holds, if it holds a nested list
 *     // The result is undefined if this NestedInteger holds a single integer
 *     const vector<NestedInteger> &getList() const;
 * };
 */

class NestedIterator {
private:
    list<NestedInteger> mylist;

public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        mylist = list<NestedInteger>(nestedList.begin(), nestedList.end());
    }
    
    int next() {
        auto tmp = mylist.front().getInteger();
        mylist.pop_front();
        return tmp;
    }
    
    bool hasNext() {
        while (!mylist.empty() && !mylist.front().isInteger())
        {
            auto first = mylist.front().getList();
            mylist.pop_front();

            for (int i = first.size() - 1; i >= 0; i--)
                mylist.insert(mylist.begin(), first[i]);
        }
        return !mylist.empty();
    }
};

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i(nestedList);
 * while (i.hasNext()) cout << i.next();
 */

解法：栈

我们可以用一个栈来代替方法一中的递归过程

具体来说，用一个栈来维护深度优先搜索时，从根节点到当前节点路径上的所有节点。由于非叶节点对应的是一个列表，我们在栈中存储的是指向列表当前遍历的元素的指针（下标）。每次向下搜索时，取出列表的当前指针指向的元素并将其入栈，同时将该指针向后移动一位。如此反复直到找到一个整数。循环时若栈顶指针指向了列表末尾，则将其从栈顶弹出

下面的代码中，$\text{C++}$ 的栈存储的是迭代器，迭代器可以起到指向元素的指针的效果

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * class NestedInteger {
 *   public:
 *     // Return true if this NestedInteger holds a single integer, rather than a nested list.
 *     bool isInteger() const;
 *
 *     // Return the single integer that this NestedInteger holds, if it holds a single integer
 *     // The result is undefined if this NestedInteger holds a nested list
 *     int getInteger() const;
 *
 *     // Return the nested list that this NestedInteger holds, if it holds a nested list
 *     // The result is undefined if this NestedInteger holds a single integer
 *     const vector<NestedInteger> &getList() const;
 * };
 */

class NestedIterator {
private:
    // pair 中存储的是列表的当前遍历位置，以及一个尾后迭代器用于判断是否遍历到了列表末尾
    stack<pair<vector<NestedInteger>::iterator, vector<NestedInteger>::iterator>> stk;
public:
    NestedIterator(vector<NestedInteger> &nestedList) {
        stk.emplace(nestedList.begin(), nestedList.end());
    }
    
    int next() {
        // 由于保证调用 next 之前会调用 hasNext，直接返回栈顶列表的当前元素，然后迭代器指向下一个元素
        auto tmp = stk.top().first->getInteger();
        stk.top().first++;
        return tmp;
    }
    
    bool hasNext() {
        while (!stk.empty())
        {
            auto& p = stk.top();
            if (p.first == p.second)
            {
                // 遍历到当前列表末尾，出栈
                stk.pop();
                continue;
            }
            if (p.first->isInteger())
                return true;
            auto& lst = p.first->getList();
            p.first++;
            stk.emplace(lst.begin(), lst.end());
        }
        return false;
    }
};

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i(nestedList);
 * while (i.hasNext()) cout << i.next();
 */