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<int> > &matrix, int k) {
		auto m = (int)matrix.size();
		if (m == 0) return 0;
		auto n = (int)matrix[0].size();
		priority_queue<S, vector<S>, greater<S>> 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;
	}
};