[LeetCode in Python] 5403 (H) find the kth smallest sum of a matrix with sorted rows 有序矩阵中的第 k 个最小数组和

题目

https://leetcode-cn.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/

给你一个 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] 是一个非递减数组

解题思路

  • 设last为第1行的数字
  • 从第1行开始,last与下面逐行计算,并更新last
    • 为了运算简便,只取前k个元素进行计算
    • 简单粗暴的二重遍历,然后排序,取前k个
    • 上面这一步也可以使用最小堆来取得TopK
  • 最后返回last[k-1]

代码

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        last = mat[0]
        for row in mat[1:]:
            last = sorted([a+b for a in last[:k] for b in row[:k]])[:k]

        return last[k-1]
爱生活,爱拉风
posted @ 2020-05-03 18:57  ET民工[源自火星]  阅读(192)  评论(0编辑  收藏  举报