Kth Smallest Number in Sorted Matrix

Description:

Find the kth smallest number in at row and column sorted matrix.

Example,

Given k = 4 and a matrix:

[
[1 ,5 ,7],
[3 ,7 ,8],
[4 ,8 ,9],
]

return 5

Challenge

O(k log n), n is the maximal number in width and height.

Solution:

struct S {
int val;
int row;
int col;
S(int v, int r, int c) : val(v), row(r), col(c) {};
};

bool operator>(S s1, S s2) {
return s1.val > s2.val;
}

class Solution {
public:
/**
* @param matrix: a matrix of integers
* @param k: an integer
* @return: the kth smallest number in the matrix
*/
int kthSmallest(vector<vector > &matrix, int k) {
auto m = (int)matrix.size();
if (m == 0) return 0;
auto n = (int)matrix[0].size();
priority_queue<S, vector, greater> q;
for (int i = 0; i < n; ++i) q.push(S(matrix[0][i], 0, i));
for (int i = 0; i < k-1; ++i) {
auto t = q.top();
q.pop();
if (t.row != m-1)
q.push(S(matrix[t.row+1][t.col], t.row+1, t.col));
}
return q.top().val;
}
};