C++ 9 lines follow up with O(k) space


  • 2

    Here's a 8 lines solution for the original question. Since we only saved iterators, space is O(1):

    class ZigzagIterator {
    private:
        queue<pair<vector<int>::iterator, vector<int>::iterator>>q;         // use a queue to record iterators
        
    public:
        ZigzagIterator(vector<int>& v1, vector<int>& v2) {
            if (v1.size()) { q.push(make_pair(v1.begin(), v1.end())); }
            if (v2.size()) { q.push(make_pair(v2.begin(), v2.end())); }
        }
    
        int next() {
            auto p = q.front(); q.pop();
            int val = *(p.first++);                                         // get the next value
            if (p.first != p.second) { q.push(p); }                         // only put next iterator back to queue if it's valid
            return val;
        }
    
        bool hasNext() {
            return q.size();
        }
    };
    

    For the follow up, idea is the same as above - just change v1/v2 to k vectors:

    class ZigzagIterator {
    private:
        queue<pair<vector<int>::iterator, vector<int>::iterator>>q;
        
    public:
        ZigzagIterator(vector<vector<int>>& vecs) {                            // input is k vectors
            for (auto v : vecs) {
                if (v.size()) { q.push(make_pair(v.begin(), v.end())); }        // only record iterators
            }
        }
    
        int next() {
            auto p = q.front(); q.pop();
            int val = *(p.first++);
            if (p.first != p.second) { q.push(p); }
            return val;
        }
    
        bool hasNext() {
            return q.size();
        }
    };
    

  • 0

    @soamaaazing Great strategy to use a queue to handle vector switching. Intuitively, one would employ an additional "indicator" to reference the current vector for next value retrieval (I used an iterator pointing to current vector for k vector case). However, the pop and push operations of queue itself can perform the job for free.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.