DeepSeek LeetCode 1439. 有序矩阵中的第 k 个最小数组和 public int kthSmallest(int[][] mat, int k)

这道题是1434的下一题，要求：从矩阵的每一行选一个数，得到数组和，找出第 k 小的和。

核心思路

多路归并 + 最小堆

• 第一行每个数作为初始和
• 依次合并下一行：当前所有和 + 新行的每个元素
• 用最小堆维护前 k 小的和

代码实现 AC

classSolution{publicintkthSmallest(int[][]mat,intk){intm=mat.length;// 行数intn=mat[0].length;// 列数// 第一行的所有数作为初始候选和PriorityQueue<Integer>pq=newPriorityQueue<>();for(intnum:mat[0]){pq.offer(num);}// 依次合并后续每一行for(introw=1;row<m;row++){PriorityQueue<Integer>nextPq=newPriorityQueue<>();// 当前所有候选和 + 新行的每个元素while(!pq.isEmpty()){intcurrSum=pq.poll();for(intnum:mat[row]){nextPq.offer(currSum+num);}}// 只保留前 k 个最小的for(inti=0;i<k&&!nextPq.isEmpty();i++){pq.offer(nextPq.poll());}}// 第 k 个最小的在堆顶（因为最后只剩 k 个）intresult=0;for(inti=0;i<k;i++){result=pq.poll();}returnresult;}}

复杂度分析

• 时间复杂度：O(m × n × k × log k)
• 空间复杂度：O(k)

更优解法：二分答案

由于 k 可能很大（最多 200），上述方法在 k 接近 n^m 时会超时，可以用二分 + 计数优化：

classSolution{publicintkthSmallest(int[][]mat,intk){intm=mat.length;intleft=0,right=0;// 最小和：每行第一个数之和for(inti=0;i<m;i++){left+=mat[i][0];right+=mat[i][n-1];}intbase=left;// 基准和（最小和）// 二分查找第 k 小的数组和while(left<right){intmid=left+(right-left)/2;// 统计有多少种数组和 <= midintcount=countArraysSumLessThan(mat,0,mid,base,k);if(count>=k){right=mid;}else{left=mid+1;}}returnleft;}privateintcountArraysSumLessThan(int[][]mat,introw,inttarget,intsum,intk){if(row==mat.length)return1;if(sum>target)return0;intcount=0;for(intval:mat[row]){count+=countArraysSumLessThan(mat,row+1,target,sum+val,k);if(count>k)break;// 剪枝}returncount;}}

两种方法对比

关键点

1. 初始堆用第一行所有数
2. 每合并一行只保留前 k 个最小和
3. 最后堆顶就是第 k 个（如果用堆弹出 k 次）