O(klogk) C++ solution with priority_queue, clean code


  • 1
    class Solution {
    public:
        int kthSmallest(vector<vector<int>>& matrix, int k) {
            int m = matrix.size(), n = matrix[0].size();
    
            auto cmp = [&matrix] (const pair<int, int>& p1, const pair<int, int>& p2) {
                return matrix[p1.first][p1.second] > matrix[p2.first][p2.second];
            };
            priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> q(cmp);
    
            q.emplace(0, 0);
            int result;
            while (k-- && !q.empty()) {
                auto t = q.top(); q.pop();
                result = matrix[t.first][t.second];
                if (t.first+1 < m) 
                    q.emplace(t.first+1, t.second);
                if (t.first == 0 && t.second+1 < n) 
                    q.emplace(t.first, t.second+1);
            }
            return result;
        }
    };
    

    It's almost the same with the No. 373 - Find K Pairs With Smallest Sums.


  • 0
    R

    @vesion are you sure it's k log k? it looks like bfs traversal in a way, where you have to push more elements the further you traverse from the starting point, which is more than k amount, and each one of the inserts taking about log(k) time. Anyways, way yours was the easiest to understand, thanks!


Log in to reply
 

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