A concise O(klogk) C++ solution

  • 0

    Inspired by the min-heap based solution of @YUANGAO1023 (java solution), I figured out a C++ one with some changes. Instead of putting the complete 1st row into the heap intially, we just need to put the 1st element.

    The key observation is that: after we fetch the top element (the minimum) from the min-heap, the next minimum of the matrix can only be found among the remaining elements of the heap + the two neighbors of the fetched top element (rightward and downward). However, we need to keep an eye on whether the two neighbors have already been put in the heap, which requires a hash set for membership test.

    struct Coordinate {
    	std::size_t x;
    	std::size_t y;  
    class Solution {
    	int kthSmallest(vector<vector<int>>& matrix, int k) {
    		// define the priority queue and the hash set 
    		auto comparer = [&matrix](const Coordinate& e1, const Coordinate& e2) {return matrix[e1.x][e1.y] > matrix[e2.x][e2.y];};
    		std::priority_queue<Coordinate, std::vector<Coordinate>, decltype(comparer)> pq(comparer); // a min-heap
    		auto equaller = [](const Coordinate& e1, const Coordinate& e2) {return e1.x == e2.x && e1.y == e2.y;};
    		auto hasher = [](const Coordinate& e) {return hash<size_t>()(e.x) ^ hash<size_t>()(e.y);};
    		auto n = matrix.size();
    		std::unordered_set<Coordinate, decltype(hasher), decltype(equaller)> set(1024 + n/10, hasher, equaller);
    		pq.push({ 0, 0 });
    		set.insert({ 0, 0 });
    		for (int i = 0; i < k - 1; i++) {
    			auto cm = pq.top();	// current minimum
    			if (cm.y + 1 < n && set.insert({ cm.x, cm.y + 1 }).second)
    				pq.push({ cm.x, cm.y + 1});
    			if (cm.x + 1 < n && set.insert({ cm.x + 1, cm.y }).second)
    				pq.push({ cm.x + 1, cm.y});
    		auto cm = pq.top(); // the kth minimum
    		return matrix[cm.x][cm.y];

    There are k times pq.pop() and at most 2k times pq.push() and set.insert(). The time complexity is at most O(klogk).

Log in to reply

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