LeetCode 378 - 有序矩阵中第 K 小的元素

在开发中，我们经常遇到需要处理大规模有序数据的场景，比如数据库分页、排行榜查询、或者处理排序过的矩阵。LeetCode 第 378 题“有序矩阵中第 K 小的元素”就是这样一个经典问题：它要求我们在一个行列都排好序的矩阵中，快速找到第 K 小的元素。本文将结合代码、算法思路和实际场景，带大家拆解这个题目的解法。

Swift社区

655人浏览 · 2025-09-12 16:16:07

Swift社区 · 2025-09-12 16:16:07 发布

在这里插入图片描述

文章目录

摘要

描述

题目要求我们在一个 n × n 的有序矩阵 中找到第 K 小的元素。矩阵的特点是：

每一行是按升序排列的；
每一列也是按升序排列的。

换句话说，矩阵不仅仅是二维数组，它在行和列两个维度上都保持有序。
我们需要返回排序后第 K 小的那个元素，而不是第 K 个不同的元素。

比如：

输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵展开后是 [1,5,9,10,11,12,13,13,15]，第 8 小的数是 13

这个问题在实际开发中也很常见，比如：

在数据库表中进行多字段排序后，查找某个偏移量的数据；
在分布式日志系统中，根据时间和序号找到某条具体日志；
在推荐系统里，找到用户兴趣排序中的第 K 个候选项。

题解答案

针对这个问题，有几个常见的思路：

最直接的做法：把矩阵“扁平化”，拉成一个数组，然后排序，再取第 K 个。
- 时间复杂度：O(n² log n²)，太大。
- 空间复杂度：O(n²)，内存开销也大。
利用小顶堆（最小堆）：每一行都是有序的，所以我们可以把每行的第一个元素丢进最小堆，每次弹出最小值，再把该元素所在行的下一个元素压入堆。这样循环 K 次，就能得到第 K 小的元素。
- 时间复杂度：O(k log n)。
二分查找 + 计数：利用矩阵整体有序的性质，在数值范围内进行二分搜索，统计小于等于 mid 的元素个数。根据计数结果收缩区间，直到找到第 K 小的数。
- 时间复杂度：O(n log(max-min))，比堆解法更节省内存。

这里我们选择 二分查找的解法，因为它更高效，也符合题目“内存复杂度优于 O(n²)”的要求。

题解代码分析

下面是 Swift 的实现代码，使用二分查找：

import Foundation

class Solution {
    func kthSmallest(_ matrix: [[Int]], _ k: Int) -> Int {
        let n = matrix.count
        var left = matrix[0][0]
        var right = matrix[n - 1][n - 1]
        
        func countLessEqual(_ mid: Int) -> Int {
            var count = 0
            var row = n - 1
            var col = 0
            
            // 从左下角开始统计 <= mid 的元素个数
            while row >= 0 && col < n {
                if matrix[row][col] <= mid {
                    count += row + 1
                    col += 1
                } else {
                    row -= 1
                }
            }
            return count
        }
        
        // 二分查找
        while left < right {
            let mid = left + (right - left) / 2
            if countLessEqual(mid) < k {
                left = mid + 1
            } else {
                right = mid
            }
        }
        
        return left
    }
}

代码解析

left 和 right 分别是矩阵中的最小值和最大值，作为二分查找的边界。
countLessEqual(mid) 用来统计矩阵中小于等于 mid 的元素个数。这里从 左下角 开始走，效率高。
- 如果当前元素 ≤ mid，那么这一列到当前行的所有数都 ≤ mid，所以 count += row + 1。
- 如果当前元素 > mid，就往上一行走。
在二分查找过程中，如果统计结果 < k，说明 mid 太小，需要右移；否则收缩右边界。
最终 left 就是第 K 小的数。

示例测试及结果

let solution = Solution()

let matrix1 = [[1,5,9],[10,11,13],[12,13,15]]
print(solution.kthSmallest(matrix1, 8))  
// 输出: 13

let matrix2 = [[-5]]
print(solution.kthSmallest(matrix2, 1))  
// 输出: -5