Binary Search - 378. Kth Smallest Element in a Sorted Matrix

79

378.Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,

return 13.

Note:
You may assume k is always valid, 1 ≤ k ≤ n2.

思路:

题目意思是在一个从左到右有序,从上到下有序的矩阵当中找到第 K 小的数,这道题目可以用堆来做,不过题目说的很明显,提示矩阵有序,那就可以考虑使用二分法来做,就是通过 mid 定位到从(0,0)到 mid 的总数和 k 来做二分,如果说计算的总数比 k 小,那么就往后搜索,否则就往前搜索。

代码:

go:

func kthSmallest(matrix [][]int, k int) int {
    if matrix == nil || len(matrix) == 0 {
        return -1
    }
    
    m := len(matrix)
    n := len(matrix[0])
    
    l := matrix[0][0]
    r := matrix[m-1][n-1]
    
    for l < r {
        mid := l + (r - l) / 2
        if countGreater(matrix, mid) >= k {
            r = mid
        } else {
            l = mid + 1 
        }
    }
    
    return l
}

func countGreater(a [][]int, k int) int {
    n := len(a)
    count := 0
    i := 0
    j := n - 1
    for i < n && j >= 0 {
        if a[i][j] > k {
            j--  
        } else {
            count += j + 1  // +1 because it's zero based indexing
            i++
        }
    }
    
    return count
}