Hard
题目描述
给你一个 m x n 的矩阵 mat，矩阵中的每一行都按非递减顺序排列，另给你一个整数 k。
你可以从每一行中选择一个元素来构成一个数组。
返回所有可能数组中第 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。
提示：
m == mat.lengthn == mat.length[i]1 <= m, n <= 401 <= mat[i][j] <= 50001 <= k <= min(200, nm)mat[i]是一个非递减数组
解题思路
这道题要求在矩阵中每行选择一个元素，使得所有可能的数组和中第k小的值。我们可以使用优先队列（最小堆）来解决这个问题。
核心思路：
- 使用最小堆来维护当前最小的数组和
- 初始状态：每行都选择第一个元素（最小值），计算初始和
- 每次从堆中取出最小和，然后生成下一批候选：对于当前选择的索引组合，尝试将每一行的索引向右移动一位
- 为了避免重复，使用集合来记录已经访问过的索引组合
具体步骤：
- 初始化堆，放入所有元素都选择第0列的组合
- 重复k次：弹出最小和，将其对应的"下一个"组合加入堆中
- “下一个"组合是指将某一行的列索引增加1得到的新组合
优化点：
- 由于每行已经排序，我们可以确保生成的候选组合是有序的
- 使用集合避免重复访问同一组合
- 时间复杂度主要取决于k和堆操作
代码实现
class Solution {
public:
int kthSmallest(vector<vector<int>>& mat, int k) {
int m = mat.size(), n = mat[0].size();
// 优先队列存储 (sum, indices)
priority_queue<pair<int, vector<int>>, vector<pair<int, vector<int>>>, greater<pair<int, vector<int>>>> pq;
set<vector<int>> visited;
// 初始状态：每行都选第一个元素
vector<int> initial(m, 0);
int initialSum = 0;
for (int i = 0; i < m; i++) {
initialSum += mat[i][0];
}
pq.push({initialSum, initial});
visited.insert(initial);
// 取k次最小值
for (int cnt = 0; cnt < k; cnt++) {
auto [currentSum, indices] = pq.top();
pq.pop();
if (cnt == k - 1) {
return currentSum;
}
// 生成下一批候选
for (int i = 0; i < m; i++) {
if (indices[i] + 1 < n) {
vector<int> newIndices = indices;
newIndices[i]++;
if (visited.find(newIndices) == visited.end()) {
visited.insert(newIndices);
int newSum = currentSum - mat[i][indices[i]] + mat[i][newIndices[i]];
pq.push({newSum, newIndices});
}
}
}
}
return -1;
}
};
class Solution:
def kthSmallest(self, mat: List[List[int]], k: int) -> int:
import heapq
m, n = len(mat), len(mat[0])
# 优先队列存储 (sum, indices_tuple)
pq = []
visited = set()
# 初始状态：每行都选第一个元素
initial_indices = tuple([0] * m)
initial_sum = sum(mat[i][0] for i in range(m))
heapq.heappush(pq, (initial_sum, initial_indices))
visited.add(initial_indices)
# 取k次最小值
for cnt in range(k):
current_sum, indices = heapq.heappop(pq)
if cnt == k - 1:
return current_sum
# 生成下一批候选
for i in range(m):
if indices[i] + 1 < n:
new_indices = list(indices)
new_indices[i] += 1
new_indices_tuple = tuple(new_indices)
if new_indices_tuple not in visited:
visited.add(new_indices_tuple)
new_sum = current_sum - mat[i][indices[i]] + mat[i][new_indices[i]]
heapq.heappush(pq, (new_sum, new_indices_tuple))
return -1
public class Solution {
public int KthSmallest(int[][] mat, int k) {
int m = mat.Length, n = mat[0].Length;
var pq = new SortedSet<(int sum, string indices)>();
var visited = new HashSet<string>();
// 初始状态：每行都选第一个元素
var initial = new int[m];
for (int i = 0; i < m; i++) initial[i] = 0;
int initialSum = 0;
for (int i = 0; i < m; i++) {
initialSum += mat[i][0];
}
string initialKey = string.Join(",", initial);
pq.Add((initialSum, initialKey));
visited.Add(initialKey);
// 取k次最小值
for (int cnt = 0; cnt < k; cnt++) {
var (currentSum, indicesStr) = pq.Min;
pq.Remove(pq.Min);
if (cnt == k - 1) {
return currentSum;
}
var indices = indicesStr.Split(',').Select(int.Parse).ToArray();
// 生成下一批候选
for (int i = 0; i < m; i++) {
if (indices[i] + 1 < n) {
var newIndices = (int[])indices.Clone();
newIndices[i]++;
string newKey = string.Join(",", newIndices);
if (!visited.Contains(newKey)) {
visited.Add(newKey);
int newSum = currentSum - mat[i][indices[i]] + mat[i][newIndices[i]];
pq.Add((newSum, newKey));
}
}
}
}
return -1;
}
}
/**
* @param {number[][]} mat
* @param {number} k
* @return {number}
*/
var kthSmallest = function(mat, k) {
let current = [mat[0][0]];
for (let i = 1; i < mat.length; i++) {
const next = [];
const minHeap = new MinPriorityQueue({ priority: x => x.sum });
const visited = new Set();
minHeap.enqueue({ sum: current[0] + mat[i][0], idx: 0 });
visited.add(0);
while (next.length < k && !minHeap.isEmpty()) {
const { sum, idx } = minHeap.dequeue().element;
next.push(sum);
for (let j = 0; j < current.length; j++) {
const newSum = current[j] + mat[i][idx];
if (newSum === sum) {
if (idx + 1 < mat[i].length && !visited.has(idx + 1)) {
minHeap.enqueue({ sum: current[j] + mat[i][idx + 1], idx: idx + 1 });
visited.add(idx + 1);
}
break;
}
}
}
current = next;
}
return current[k - 1];
};
复杂度分析
| 复杂度类型 | 复杂度 | 说明 |
|---|---|---|
| 时间复杂度 | O(k × m × log(k × m)) | 最多处理k×m个状态，每次堆操作为O(log(k×m)) |
| 空间复杂度 | O(k × m) | 优先队列和visited集合的空间开销 |