C++ Clean Priority Queue (14 lines)


  • 0

    Exact same idea as https://discuss.leetcode.com/topic/52948/share-my-thoughts-and-clean-java-code.

    class Solution {
    public:
        using element = tuple<int, int, int>; // (value, row, col)
        
        int kthSmallest(vector<vector<int>>& M, int k) {
             // N = number of rows/cols
            int N = M.size();
            
            // Set up priority queue, which records each element's position
            auto cmp = [](element a, element b) { return get<0>(a) > get<0>(b); };
            priority_queue<element, vector<element>, decltype(cmp)> pq(cmp);
            
            // Push all elements in first row
            for (int col = 0; col < N; col++) pq.emplace(M[0][col], 0, col);
            
            // Repeat k times:
            // 1. Pop min element from pq, which is at M[popRow][popCol]
            // 2. Push element M[popCol+1][popRow] if it's valid
            element res;
            for (int i = 0; i < k; i++) {
                res = pq.top(); pq.pop();
                int popRow = get<1>(res);
                int popCol = get<2>(res);
                if (popRow < N-1) pq.emplace(M[popRow+1][popCol], popRow+1, popCol);
            }
            return get<0>(res);
        }
    };
    

Log in to reply
 

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