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

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

1. 问题描述

2. 思路

1. 左上角元素作为left，右下角元素作为right
2. 执行二分查找

3. 代码

func kthSmallest(matrix [][]int, k int) int {
    n := len(matrix)
    left, right := matrix[0][0], matrix[n-1][n-1]   
    
    for left <= right {
        mid := left + (right - left) / 2
        // 计算matrix矩阵中小于mid 的元素个数
        smallNum := countSmall(matrix, mid, n)
        // fmt.Println(smallNum)
        if smallNum == k {
            right = mid - 1  // 不要怕错过，最终一定会在此处汇聚
        } else if smallNum < k {
            // mid太小，增大mid
            left = mid + 1
        } else {
            right = mid - 1
        }
    }
    return left
}

// 计算matrix矩阵中小于mid 的元素个数
func countSmall(matrix [][]int, mid int, n int) int {
    var num int
    i, j := n - 1, 0

    for i >= 0 && j <= n - 1 {
        if matrix[i][j] <= mid {
            j++
            num += i + 1
        } else {
            i--
        }
    }
    return num
}