- Is
kalways valid?
Input: [
[ 1, 2],
[ 3, 5],
],
k = 2
Output: 2kis the smallest number or the largest number in the matrix.
Input: [
[ 1, 2],
[ 3, 5],
],
k = 1 or 4
Output: 1 or 5- There are some duplicates that is
kth smallest number or in the same row or column.
Input: [
[ 1, 2],
[ 1, 3],
],
k = 2
Output: 1
Input: [
[ 1, 2],
[ 3, 3],
],
k = 3
Output: 3
Input: [
[-5, -4],
[-5, -4],
],
k = 2fun kthSmallest(matrix: Array<IntArray>, k: Int): Int {
val maxHeap = PriorityQueue<Int>() { n1, n2 -> n2 - n1 }
for (r in 0 until matrix.size) {
for (c in 0 until matrix[r].size) {
maxHeap.add(matrix[r][c])
if (maxHeap.size > k) maxHeap.poll()
}
}
return maxHeap.peek()!!
}- Time Complexity: Iterating all item in the matrix takes
O(n^2), andadd()/polltakesO(lg k), total takesO(n^2 lg k). - Space Complexity:
O(k)for heap, but the worst case isk = n^2, so it'sO(n^2)which does not meet the requirement.
We can treat the matrix as n sorted linked list, and we can use a min heap to store the smallest element of each list. We can add the first element of each row to the heap, and then we can poll the smallest element from the heap and add the next element of the same row to the heap. We repeat this process k times.
1 3 6 => 1 -> 3 -> 6
p1
2 5 8 => 2 -> 5 -> 8
p2
4 7 9 => 4 -> 7 -> 9
p3fun kthSmallest(matrix: Array<IntArray>, k: Int): Int {
val minHeap = PriorityQueue<Pair<Int, Int>>() { p1, p2 ->
matrix[p1.first][p1.second] - matrix[p2.first][p2.second]
}
for (row in 0 until matrix.size) {
minHeap.add(row to 0)
}
var kk = k
// Poll k - 1 times
while (1 < kk && minHeap.isNotEmpty()) {
val pair = minHeap.poll()
if (pair.second + 1 < matrix[0].size) {
minHeap.add(pair.first to (pair.second + 1))
}
kk--
}
// The kth smallest number
val peek = minHeap.peek()
return matrix[peek.first][peek.second]
}The matrix rows and columns are sorted, so the minimum must be A[0][0] and the maximum must be at A[n-1][n-1]:
Monotonicity: The Kth smallest elements exists in the range A[0][0]..A[n-1][n-1], given a number middle in the range left..right, then the count of all numbers <= middle will be the left portion of the matrix. It will be within the range that are covered during search.
We can set left and right to be the minimum and maximum value in the matrix, and we can apply binary search to find the Kth smallest number in the matrix. It's similiar to 240. Search a 2D Matrix II but the search space is reduced by the count of numbers that are <= middle.
- Search Space Reduction: Actually, the problem can be converted to search the first number that meets the specific condition (The pattern is similar to 278. First Bad Version), condition is
count(<=X) is k.
[X, X, X, O, O, O, O]
^ The first number that has count(<=x) is `k`
We can count how many numbers k in the matrix, we can reduce the search space by this count. If the count < k, then we search the right part, otherwise, we search the left part. Here we don't return the value immediately when count == k, because the value might not be in the matrix, it just matches the count, so we keep shrinking the search space until breaking the while loop.
- Result Guarantee: The Kth smallest number must be in
left..right, and we keep reducing the search space by the conditionthe count <= kuntilleft == right, then the value is what we're looking for.
Why does the Kth smallest number is the value in the matrix when breaking the loop?
if count<k , meaning that you need to make mid become bigger.
else, meaning that you should limit the high.
then you get low == high and k == count.
"find the first occurrence of an element".
How lo will always be an element in the matrix array in the end ? In the end, lo == hi. lo is the smallest number satisfying that count(lo)>=k. **let's assume that lo doesn't exist in matrix.** We must have the biggest num in matrix that is less than lo. Because ans didn't exist in matrix, so count(num) >= k also holds which contradicts our previous assumption that lo is the smallest. So lo does exist in matrix.
Based on the above characteristics, we can apply binary search: We set left and right to the minimum and maximum first, we calculate the middle based on left and right, and we count the item numbers that is <= middle.
* If count < k, the range is not covered k, we search right part.
* Otherwise (count >= k), we search left part.
答案的上限是
matrix[0][0],下限是matrix[N-1][N-1]．对于这个区间的任意一个ｘ，我们可以计算出这个矩阵里小于等于ｘ的元素有多少，定义为smallerOrEqual(x)，如果smallerOrEqual(x)<k，说明这个ｋ太小了不是想要的答案，应该往上调整. 反之smallerOrEqual(x)>=k，说明ｋ可能是答案，但可以再缩小一点试一试。
count(≤middle) ≥ k表示可能是解，但有可能太大，必要但是不充分。二分夾擠最後夾出第一個滿足count(<=middle) ≤ k，再小一點就不會滿足這條件
fun kthSmallest(matrix: Array<IntArray>, k: Int): Int {
val n = matrix.size
var left = matrix[0][0]
var right = matrix[n - 1][n - 1]
var result = 0
while (left <= right) {
val middle = left + (right - left) / 2
val count = countLessOrEqual(matrix, middle)
/**
* We won't return immediately if count == k, because it the middle might not be the value in the matrix, it's just count match.
* Or see failed case:
* [
* [1, 2]
* [1, 3]
* ]
* k = 1
* if (count == k) return middle
*/
// But we can update the result when (count == k)
if (count == k) result = middle
// |------------------| matrix
// k
// |------| count
// count is too small that the search range doesn't cover the kth smallest number.
if (count < k) {
// search right part
left = middle + 1
} else {
// k <= count, that means the current search range covers the kth smallest number,
// so we have to update the result, and search left part
right = middle - 1
}
}
return result
}This lead to WA:
if (count < k) {
left = middle + 1
} else if (count == k) {
result = middle
right = middle - 1
} else if (k < count) {
right = middle - 1
}Failed at the test case:
[
[1,2]
[3,3]
]
k = 3
Or equivalently, we don't update the answer when count == k, we try to find the first number that meets the condition countLessOrEqual() == k, then we can return the left pointer when while loop breaks.
fun kthSmallest(matrix: Array<IntArray>, k: Int): Int {
val n = matrix.size
var left = matrix[0][0]
var right = matrix[n - 1][n - 1]
var result = left
while (left <= right) {
val middle = left + (right - left) / 2
val count = countLessOrEqual(matrix, middle)
if (k < count) {
// Search left part
right = middle - 1
} else if (count < k){
// Search right part
left = middle + 1
} else if (count == k) { // We are not sure if it's the element element that count == k, so we keep searching the left part.
right = middle - 1
}
}
return left
}
private fun countLessOrEqual(matrix: Array<IntArray>, middle: Int): Int {
var count = 0
var row = matrix.size - 1
for (col in 0 until matrix.size) {
while (row >= 0 && matrix[row][col] > middle) {
row--
}
count += row + 1
}
return count
}1 3 6
2 5 8
4 7 9
k = 4
Left 1 1 3 3 4
Right 9 5 5 4 4
Middle 5 3 4 3 4
Count 5 3 4 3 4// Input
1, 5, 9
10, 11, 13
12, 13, 15
k = 8
// Count
L R M C
1, 15, 8, 2
9, 15, 12, 6
13, 15, 14, 8
// It will return 14 immediately when count == 8, it's WA, 14 is not in our matrix
// Correct way:
L R M C
1, 15, 8, 2
9, 15, 12, 6
13, 15, 14, 8
13, 13, 13, 8
// We shrink the search range to be left == right, then the value is what we're looking for.