[LeetCode] 251、 展开二维向量

题目描述

请设计并实现一个能够展开二维向量的迭代器。该迭代器需要支持 next 和 hasNext 两种操作。

示例：

Vector2D iterator = new Vector2D([[1,2],[3],[4]]);

iterator.next(); // 返回 1
iterator.next(); // 返回 2
iterator.next(); // 返回 3
iterator.hasNext(); // 返回 true
iterator.hasNext(); // 返回 true
iterator.next(); // 返回 4
iterator.hasNext(); // 返回 false

参考代码

class Vector2D {
public:
    vector<vector<int>>::iterator outerIter;
    vector<vector<int>>::iterator outerEnd;
    vector<int>::iterator innerIter;
    
    Vector2D(vector<vector<int>>& v) {
        //因为第一个next要返回第一个元素，所以只能在一开始添加一个工具dummy元素
        v.insert(v.begin(), vector<int>(1,0));
        outerEnd = v.end();
        outerIter = v.begin();
        innerIter = v[0].begin();
    }
    
    int next() {
        ++innerIter;
        //直到找到一个非空vector为止
        while(innerIter == outerIter->end()){
            ++outerIter;
            innerIter = outerIter->begin();
        }
        return (*innerIter);
    }
    
    bool hasNext() {
        auto tempInner = innerIter;
        auto tempOuter = outerIter;
        
        ++tempInner;
        if(tempInner != outerIter->end()) return true; //如果直接在第二维找到
        
        ++tempOuter;
        if(tempOuter == outerEnd) return false; //第一维已经搜索完还没有搜索到，失败
        tempInner = tempOuter->begin();
        
        while(tempInner == tempOuter->end()){ //如果第二维的vector为空，持续搜索
            ++tempOuter;
            if(tempOuter == outerEnd) return false;
            tempInner = tempOuter->begin();
        }
        
        if(tempInner != tempOuter->end()) return true; //最后检查第二维的vector是不是空的
        return true;
    }
};