更新 5. 有序矩阵的 Kth Element 二分查找法实现

原本实现,在计算 matrix[i][j] <= mid 的 count 时,耗时较高
This commit is contained in:
lichaoxi 2020-01-01 10:30:04 +08:00 committed by GitHub
parent 5f005e2263
commit 205a1532da
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -146,20 +146,26 @@ return 13.
```java ```java
public int kthSmallest(int[][] matrix, int k) { public int kthSmallest(int[][] matrix, int k) {
int m = matrix.length, n = matrix[0].length; int n = matrix.length;
int lo = matrix[0][0], hi = matrix[m - 1][n - 1]; int low = matrix[0][0], high = matrix[n-1][n-1];
while (lo <= hi) { while (low < high) {
int mid = lo + (hi - lo) / 2; int mid = (low + high) >> 1;
int cnt = 0; int count = 0;
for (int i = 0; i < m; i++) { for (int i = 0, j = n - 1; i < n && j >= 0;) {
for (int j = 0; j < n && matrix[i][j] <= mid; j++) { if (matrix[i][j] <= mid) {
cnt++; count += j + 1;
i++;
} else {
j--;
} }
} }
if (cnt < k) lo = mid + 1; if (count < k) {
else hi = mid - 1; low = mid + 1;
} else {
high = mid;
}
} }
return lo; return high;
} }
``` ```