//Best first search, time complexity O(klog(n)). The size of heap is at most
2*n. Every time an element is popped from the heap or inserted to the heap,
O(log(n)). We need to do the pop and insert operations k times. So the time
complexity is O(klog(n)).
//Space complexity O(n), because there are at most 2*n nodes in the priority
queue.
struct node
{
int index;
int val;
node(int index_, int val_): index(index_), val(val_){}
bool operator > (const node& other) const
{
return val > other.val;
}
};
static int findKthSmallest2D(const vector> &mat, int k)
{
int res = INT_MIN;
int m = mat.size();
if (m == 0)
{
return res;
}
int n = mat[0].size();
if (n == 0)
{
return res;
}
priority_queue, greater> minHeap;
unordered_set hash; // to avoid duplication.
minHeap.push(node(0, mat[0][0]));
hash.insert(0);
while (!minHeap.empty())
{
node top = minHeap.top();
minHeap.pop();
k--;
if (k == 0)
{
return top.val;
}
int i = top.index / n;
int j = top.index % n;
int c1 = (i + 1) * n + j;
int c2 = i * n + j + 1;
if (i + 1 < m && hash.find(c1) == hash.end())
{
minHeap.push(node(c1, mat[i + 1][j]));
hash.insert(c1);
}
if (j + 1 < n && hash.find(c2) == hash.end())
{
minHeap.push(node(c2, mat[i][j + 1]));
hash.insert(c2);
}
}
return res;
}
