给你一个 m * n 的矩阵 mat,以及一个整数 k ,矩阵中的每一行都以非递减的顺序排列。 你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。
示例 :
输入: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 。
解法1: 暴力求解 链接:https://leetcode-cn.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/jian-ji-de-python3jie-fa-by-please_ac/
def kthSmallest(self, mat: List[List[int]], k: int) -> int: n = len(mat) pre = [0] for i in range(n): cur = [a + b for a, b in itertools.product(pre, mat[i])] pre = heapq.nsmallest(min(k, len(cur)), cur) return pre[-1]解法2: 二分法 链接:https://leetcode-cn.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/er-fen-by-newbie-19-3/
class Solution { public: vector<vector<int>>temp; int m,n; int kthSmallest(vector<vector<int>>& mat, int k) { temp=mat; m=mat.size(),n=mat[0].size(); int left=0,right=0; for(int i=0;i<m;i++) left+=mat[i][0],right+=mat[i].back(); int init=left; while(left<right){ int mid=(left+right)>>1; int num=1; dfs(mid,0,init,num,k); if(num>=k) right=mid; else left=mid+1; } return left; } void dfs(int mid,int index,int sum,int& num,int k){ if(sum>mid||index==m||num>k) return; dfs(mid,index+1,sum,num,k); for(int i=1;i<n;i++){ if(sum+temp[index][i]-temp[index][0]<=mid){ num++; dfs(mid,index+1,sum+temp[index][i]-temp[index][0],num,k); }else{ break; } } } };解法3: 小顶堆 链接:https://leetcode-cn.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/bao-li-jie-fa-zui-xiao-dui-by-coldme-2/
import heapq class Solution: def kthSmallest(self, mat, k: int) -> int: m, n = len(mat), len(mat[0]) # 初始化指针 pointers = [0] * m # 初始化heap heap = [] curr_sum = 0 for i in range(m): curr_sum += mat[i][0] heapq.heappush(heap, [curr_sum, tuple(pointers)]) # 初始化seen seen = set() seen.add(tuple(pointers)) # 执行k次 for _ in range(k): # 从堆中pop出curr_sum(最小数组和)和pointers(指针数组) curr_sum, pointers = heapq.heappop(heap) # 每个指针轮流后移一位,将new_sum(新的数组和)和new_pointers(新的指针数组)push入堆 for i, j in enumerate(pointers): if j < n - 1: new_pointers = list(pointers) new_pointers[i] = j + 1 new_pointers = tuple(new_pointers) if new_pointers not in seen: new_sum = curr_sum + mat[i][j + 1]- mat[i][j] heapq.heappush(heap, [new_sum, new_pointers]) seen.add(new_pointers) return curr_sum