#378. 有序矩阵中第K小的元素

给定一个 n * n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。

示例:

1
2
3
4
5
6
7
8
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,

返回 13。

提示:
你可以假设 k 的值永远是有效的,1 ≤ k ≤ n^2

#写在前面

写这个的根本目的是记录一些之前根本不会注意到的小错误。在我写完去看官方题解的时候,官方的求中位数方法是:

1
int mid = left + ((right - left) >> 1);

自己的是:

1
int temp = (min + max) / 2;

感叹自己不会移位但却少了一步减法的时候,却没有考虑到两个 int 相加是可能会溢出的。

所以今日小tip:

  • 两个 int 相加是可能溢出的。

#二分查找

最简单的,取出来排序,这个就不说了,肯定不是最优雅的解法。

老实说我最开始的想法是从左上角开始向下向右数个数,维护两个队列,一个向右的,一个向下的。但我意识到完全存在着一个数的向下可能小于前面某一个数的向右。完全没有办法按序输出。此思路作罢。

第二个思路是 i * i 的值如果大于 k,那么结果一定在 i * i的子矩阵中。但除了 matrix[i][i]这个点之外,所有其它的点都有可能是结果,所以没有办法通过这个条件继续缩小范围了。

想了半小时没有想到好的解法,我选择了去看官方题解,简单了解了一下思路后,回来自己想。

#判断特殊情况

已经养成习惯了,在每个方法的最开始将特殊情况判断一下。

1
2
3
4
5
6
7
8
9
class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
}
}

#进行一次查找

先写一次查找的结果,看看对不对。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
int min = matrix[0][0]; // 初始化最小值
int max = matrix[n - 1][n - 1]; // 初始化最大值
int i = n - 1, j = 0; // 从左下角开始找
int count = 0; // 记录个数
int temp = (min+max)/2; // 中位数开始找
while (i >= 0 && j < n) { // 超过数组范围则退出
if(matrix[i][j] <= temp){ // 向左
j+=1;
continue;
} else { // 向上
count += j; // 向上之前,把第i行前j-1个元素都算上
// 索引默认0开始,j-1+1=j
i -= 1;
}
}
count += (i+1)*n; // 如果因为j越界的原因退出,证明上面还有i+1行没有算
return count;
}
}
1 3 5 7 9 11
2 4 6 8 10 12
3 5 7 9 11 13
4 6 8 10 12 14
5 7 9 11 13 15
6 8 10 12 14 16

结果是 18,正确。

#迭代查找

接着就可以写更新策略了,我想的是,如果现在的数小于要求的数,那一定是在后面,但是如果等于却不一定,因为这个中位数可能是矩阵中不存在的数,但最终结果一定会小于等于现在的数,所以>=的策略可以写在一起。

最后的的迭代终止条件就是min==max

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
int min = matrix[0][0];
int max = matrix[n - 1][n - 1];
while (true) {
int i = n - 1, j = 0;
int count = 0;
int temp = (min + max) / 2;
while (i >= 0 && j < n) {
if (matrix[i][j] <= temp) {
j += 1;
continue;
} else {
count += j;
i -= 1;
}
}
count += (i + 1) * n;
if (count < k) {
min = temp;
} else {
max = temp;
}
if (min == max) {
break;
}
}
return min;
}
}

但是这个会死循环卡住,我用纸自己跑了一下,发现假如min = 7, max = 8,再跑一轮的时候迭代还会是min = 7, max = 8,无限循环了。解决方法是判断迭代值和 min 是否相等:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
int min = matrix[0][0];
int max = matrix[n - 1][n - 1];
while (true) {
int i = n - 1, j = 0;
int count = 0;
int temp = (min + max) / 2;
while (i >= 0 && j < n) {
if (matrix[i][j] <= temp) {
j += 1;
continue;
} else {
count += j;
i -= 1;
}
}
count += (i + 1) * n;
if (count < k) {
if (temp == min) {
min += 1;
break;
}
min = temp;
} else {
max = temp;
}
if (min == max) {
break;
}
}
return min;
}
}

通过了。当然,官方题解里的 +1 就很灵性,但是我想不到。。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
int min = matrix[0][0];
int max = matrix[n - 1][n - 1];
while (true) {
int i = n - 1, j = 0;
int count = 0;
int temp = (min + max) / 2;
while (i >= 0 && j < n) {
if (matrix[i][j] <= temp) {
j += 1;
continue;
} else {
count += j;
i -= 1;
}
}
count += (i + 1) * n;
if (count < k) {
min = temp + 1;
} else {
max = temp;
}
if (min == max) {
break;
}
}
return min;
}
}

#完整代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
package leetcode;

public class KthSmallestElementInASortedMatrix {
public static void main(String[] args) {
int[][] matrix = { { 1, 3, 5, 7, 9, 11 }, { 2, 4, 6, 8, 10, 12 }, { 3, 5, 7, 9, 11, 13 },
{ 4, 6, 8, 10, 12, 14 }, { 5, 7, 9, 11, 13, 15 }, { 6, 8, 10, 12, 14, 16 } };
int k = 18;
KthSmallestElementInASortedMatrixSolution sol = new KthSmallestElementInASortedMatrixSolution();
int res = sol.kthSmallest(matrix, k);
System.out.println(res);
}
}

class KthSmallestElementInASortedMatrixSolution {
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
if (n == 0)
return 0;
if (n == 1)
return matrix[0][0];
int min = matrix[0][0];
int max = matrix[n - 1][n - 1];
while (true) {
int i = n - 1, j = 0;
int count = 0;
int temp = (min + max) / 2;
while (i >= 0 && j < n) {
if (matrix[i][j] <= temp) {
j += 1;
continue;
} else {
count += j;
i -= 1;
}
}
count += (i + 1) * n;
if (count < k) {
if (temp == min) {
min += 1;
break;
}
min = temp;
} else {
max = temp;
}
if (min == max) {
break;
}
}
return min;
}
}
1
2
3
% javac leetcode/KthSmallestElementInASortedMatrix.java
% java leetcode.KthSmallestElementInASortedMatrix
8

完。