有序矩阵中的第 k 个最小数组和

标签: 数组 二分查找 矩阵 堆（优先队列）

难度: Hard

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。

示例 1：

输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。

示例 2：

输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17

示例 3：

输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。

示例 4：

输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12

提示：

  • m == mat.length
  • n == mat.length[i]
  • 1 <= m, n <= 40
  • 1 <= k <= min(200, n ^ m)
  • 1 <= mat[i][j] <= 5000
  • mat[i] 是一个非递减数组

Submission

运行时间: 38 ms

内存: 17.7 MB

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        m = len(mat)
        n = len(mat[0])
        val = sum(nums[0] for nums in mat)
        q = [tuple([val] + [0] * m)]

        for _ in range(k - 1):
            top = heappop(q)
            val = top[0]
            indexs = list(top[1:])

            for i, index in enumerate(indexs):
                if index + 1 < n:
                    indexs[i] += 1
                    new_val = val + mat[i][index + 1] - mat[i][index]
                    heappush(q, tuple([new_val] + indexs))
                    indexs[i] -= 1

                if index != 0:
                    break

            

        return q[0][0]

Explain

此题解采用了最小堆策略来解决问题。它首先计算出矩阵每一行的第一个元素的和作为初始状态，并且跟踪每一行中当前选中的元素的索引。解的核心在于迭代地从堆中取出当前最小的数组和，然后探索通过在某一行中前进到下一个元素而形成的新的数组和。对于每个取出的元素，尝试每行的下一个元素，如果生成了更小的数组和，则将其推入堆中。这个过程重复k次，以确保找到第k个最小的数组和。

时间复杂度: O(k * m * log k)

空间复杂度: O(k * m)

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        m = len(mat)  # 矩阵的行数
        n = len(mat[0])  # 矩阵的列数
        val = sum(nums[0] for nums in mat)  # 计算初始元素和，即每行的第一个元素的和
        q = [tuple([val] + [0] * m)]  # 初始化堆，包含初始和与每行选中的元素的索引

        for _ in range(k - 1):  # 迭代k-1次，因为第k个元素将在最后从堆中取出
            top = heappop(q)  # 从堆中取出当前最小的和
            val = top[0]
            indexs = list(top[1:])

            for i, index in enumerate(indexs):  # 尝试更新当前行的元素
                if index + 1 < n:  # 如果当前行有更多元素可以选择
                    indexs[i] += 1
                    new_val = val + mat[i][index + 1] - mat[i][index]  # 计算新的数组和
                    heappush(q, tuple([new_val] + indexs))  # 将新的和推入堆中
                    indexs[i] -= 1  # 恢复索引，为下一行的更新做准备

                if index != 0:  # 优化：只在索引非零时继续，防止重复
                    break

        return q[0][0]  # 返回堆中最小的元素，即第k个最小的数组和

Explore

在这个问题中，我们的目标是找到第k个最小的数组和。使用最小堆的优势在于其能够快速地提供当前最小的元素，这是解决这类问题的关键。每次从堆中取出元素（当前最小）后，只需要考虑少数几种可能产生下一个最小和的情况，这些新的和可以快速地插入到堆中。相比之下，虽然二叉搜索树或平衡树同样可以维护有序的元素集合，但在频繁插入和删除操作时，平衡树需要进行额外的旋转操作来保持平衡，可能会导致较高的处理开销。此外，堆结构简单，实现和操作更为直观，特别适用于需要频繁提取最小元素和更新元素的场景。

这种策略不会遗漏其他潜在的更小的数组和。在每次从堆中取出当前最小的数组和之后，我们通过考虑该数组和中每个元素的下一个可能的元素来生成新的数组和。这保证了所有可能产生更小数组和的情况都会被考虑到。只有当前最小的数组和被完全探索后，我们才会从堆中取出下一个最小的数组和。这种方法确保了每次都是全局最小的数组和被处理，从而避免了遗漏任何可能的更小数组和。

这种优化策略基于避免重复处理相同的元素组合的原理。在从堆中取出一个元素后，我们尝试通过移动到当前行的下一列来生成新的数组和。一旦我们在某一行进行了更新，后续行的相同更新就会产生重复的数组组合，因为前面的行的元素位置没有变化。因此，我们在更新某行的索引后，如果不是第一个元素（索引非零），就停止进一步的行更新操作，这样可以防止在堆中生成重复的数组和，提高算法的效率。