想了解更多数据结构以及算法题,可以关注微信公众号“数据结构和算法”,每天一题为你精彩解答。也可以扫描下面的二维码关注
给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第k小的元素。
请注意,它是排序后的第k小元素,而不是第k个元素。
示例:
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。
答案:
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[2] - b[2]);
for (int j = 0; j < n; j++)
pq.offer(new int[]{0, j, matrix[0][j]});
for (int i = 0; i < k - 1; i++) {
int[] cur = pq.poll();
if (cur[0] == n - 1)
continue;
pq.offer(new int[]{cur[0] + 1, cur[1], matrix[cur[0] + 1][cur[1]]});
}
return pq.poll()[2];
}
解析:
PriorityQueue会对添加进去的数据进行排序,其实他就是一个堆,在这里他是个最小堆,也就是最顶端的元素是最小的(虽然他是数组结构,但数组的位置是有关联的),每添加一个元素他都会往上调整,删除的时候往下调整,并且他有两个最重要的函数一个是siftDown,一个是siftUp。上面的代码中它添加的是个数组,数组的最后一个元素是我们添加的值,前两个元素是这个值在矩阵中(x,y)的坐标。他是先把矩阵中第一行的元素全部添加进去,后面再进行k-1次循环。每次循环的时候都会把最小的给移除掉,然后把它紧挨着的下一个元素添加进去。因为是最小堆,所以当我们移除了k-1次的时候,这时堆的顶端就是第k小的元素。我们还以上面的例子来画个图加深一下理解
题中矩阵的每行每列都是排过序的,所以我们还可以想到另一种方法,使用二分法查找,之前介绍过二分法查找,不会的可以看下202,查找-二分法查找
public int kthSmallest(int[][] matrix, int k) {
int low = matrix[0][0], high = matrix[matrix.length - 1][matrix[0].length - 1];
while (low < high) {
int mid = low + (high - low) / 2;
int count = 0, j = matrix[0].length - 1;
for (int i = 0; i < matrix.length; i++) {
while (j >= 0 && matrix[i][j] > mid)
j--;
count += (j + 1);
}
if (count < k)
low = mid + 1;
else
high = mid;
}
return low;
}
因为矩阵的行和列都是排过序的,这里先找到最中间的值,count表示的是比mid小的值有多少个,每次都是用每行的最右边的一个值和mid比较,直到比mid大,才会执行上面的while循环,然后再往前找,这里的第count个值是大于mid的最小值,所以我们不能在count==k的时候直接return。