import heapq
class Solution:
def kthSmallest(self, mat, k):
m, n = len(mat), len(mat[0])
initial = sum(row[0] for row in mat)
heap = [(initial, [0]*m)]
visited = set(tuple([0]*m))
for _ in range(k):
curr_sum, indices = heapq.heappop(heap)
for i in range(m):
if indices[i] + 1 < n:
nxt = list(indices)
nxt[i] += 1
nxt_tuple = tuple(nxt)
if nxt_tuple not in visited:
new_sum = curr_sum - mat[i][indices[i]] + mat[i][nxt[i]]
heapq.heappush(heap, (new_sum, nxt))
visited.add(nxt_tuple)
return curr_sum
#include <vector>
#include <queue>
#include <set>
using namespace std;
class Solution {
public:
int kthSmallest(vector<vector<int>>& mat, int k) {
int m = mat.size(), n = mat[0].size();
vector<int> idx(m, 0);
int initial = 0;
for (int i = 0; i < m; ++i) initial += mat[i][0];
typedef pair<int, vector<int>> Node;
auto cmp = [](const Node& a, const Node& b) { return a.first > b.first; };
priority_queue<Node, vector<Node>, decltype(cmp)> minHeap(cmp);
set<vector<int>> visited;
minHeap.push({initial, idx});
visited.insert(idx);
int curr_sum = 0;
for (int t = 0; t < k; ++t) {
auto [sum, indices] = minHeap.top(); minHeap.pop();
curr_sum = sum;
for (int i = 0; i < m; ++i) {
if (indices[i] + 1 < n) {
vector<int> nxt = indices;
nxt[i]++;
if (!visited.count(nxt)) {
int new_sum = sum - mat[i][indices[i]] + mat[i][nxt[i]];
minHeap.push({new_sum, nxt});
visited.insert(nxt);
}
}
}
}
return curr_sum;
}
};
import java.util.*;
class Solution {
public int kthSmallest(int[][] mat, int k) {
int m = mat.length, n = mat[0].length;
int[] indices = new int[m];
int initial = 0;
for (int i = 0; i < m; ++i) initial += mat[i][0];
PriorityQueue<Node> heap = new PriorityQueue<>();
Set<String> visited = new HashSet<>();
heap.offer(new Node(initial, indices.clone()));
visited.add(Arrays.toString(indices));
int curr_sum = 0;
for (int t = 0; t < k; ++t) {
Node node = heap.poll();
curr_sum = node.sum;
for (int i = 0; i < m; ++i) {
if (node.indices[i] + 1 < n) {
int[] nxt = node.indices.clone();
nxt[i]++;
String key = Arrays.toString(nxt);
if (!visited.contains(key)) {
int new_sum = node.sum - mat[i][node.indices[i]] + mat[i][nxt[i]];
heap.offer(new Node(new_sum, nxt));
visited.add(key);
}
}
}
}
return curr_sum;
}
static class Node implements Comparable<Node> {
int sum;
int[] indices;
Node(int s, int[] idx) {
sum = s;
indices = idx;
}
public int compareTo(Node other) {
return Integer.compare(this.sum, other.sum);
}
}
}
class MinHeap {
constructor() {
this.heap = [];
}
push(item) {
this.heap.push(item);
this._bubbleUp(this.heap.length - 1);
}
pop() {
if (this.heap.length === 1) return this.heap.pop();
const top = this.heap[0];
this.heap[0] = this.heap.pop();
this._bubbleDown(0);
return top;
}
_bubbleUp(idx) {
while (idx > 0) {
let parent = (idx - 1) >> 1;
if (this.heap[parent][0] <= this.heap[idx][0]) break;
[this.heap[parent], this.heap[idx]] = [this.heap[idx], this.heap[parent]];
idx = parent;
}
}
_bubbleDown(idx) {
const length = this.heap.length;
while (true) {
let left = 2 * idx + 1, right = 2 * idx + 2, smallest = idx;
if (left < length && this.heap[left][0] < this.heap[smallest][0]) smallest = left;
if (right < length && this.heap[right][0] < this.heap[smallest][0]) smallest = right;
if (smallest === idx) break;
[this.heap[smallest], this.heap[idx]] = [this.heap[idx], this.heap[smallest]];
idx = smallest;
}
}
size() { return this.heap.length; }
}
var kthSmallest = function(mat, k) {
const m = mat.length, n = mat[0].length;
let initial = 0;
for (let i = 0; i < m; ++i) initial += mat[i][0];
const heap = new MinHeap();
const visited = new Set();
const start = Array(m).fill(0);
heap.push([initial, start]);
visited.add(start.join(','));
let curr_sum = 0;
for (let t = 0; t < k; ++t) {
const [sum, indices] = heap.pop();
curr_sum = sum;
for (let i = 0; i < m; ++i) {
if (indices[i] + 1 < n) {
const nxt = indices.slice();
nxt[i]++;
const key = nxt.join(',');
if (!visited.has(key)) {
let new_sum = sum - mat[i][indices[i]] + mat[i][nxt[i]];
heap.push([new_sum, nxt]);
visited.add(key);
}
}
}
}
return curr_sum;
};
Given a matrix mat of size m x n, where each row of mat is sorted in non-decreasing order, you are asked to find the k-th smallest sum that can be formed by picking exactly one element from each row and summing them up.
Key constraints:
k-th smallest such sum (1-based indexing).
At first glance, the problem may seem like a variant of finding the k-th smallest element in a sorted array, but here, the "elements" are sums generated by picking one value from each row. Since each row is sorted, we can try to exploit this property.
The brute-force approach would be to generate all possible combinations (which is n^m if there are n columns and m rows), compute their sums, sort them, and pick the k-th smallest. However, this is computationally infeasible for larger matrices.
Instead, we look for a more efficient way, possibly using a min-heap (priority queue) to always expand the next smallest sum, similar to the approach used in "Kth Smallest Element in a Sorted Matrix". By always expanding the next smallest sum, we can avoid generating all combinations up front and only focus on the most promising candidates.
The solution uses a min-heap (priority queue) to efficiently generate the next smallest possible sum by picking one element from each row. Here’s a step-by-step breakdown:
k times. The sum popped from the heap on the k-th iteration is the answer.This approach is efficient because it always expands the next smallest sum and avoids recomputing or revisiting the same index combinations.
Suppose:
mat = [[1,3,11], [2,4,6], [5,7,8]]
k = 5
1 + 2 + 5 = 8 (indices [0,0,0]).
12.
This process ensures we always expand the next smallest possible sum, efficiently reaching the k-th smallest.
Brute-force approach:
O(n^m \cdot m) (where n is number of columns, m is number of rows), since all n^m combinations are generated and summed.O(n^m) for storing all sums before sorting.O(k \cdot m \cdot \log k), since at most k elements are pushed/popped from the heap, and for each, up to m new combinations are tried.O(k) for the heap and the visited set.
The heap-based approach is efficient even for reasonably large k and matrix sizes because it avoids expanding all combinations.
This problem is a classic example of using a priority queue (min-heap) to efficiently generate and expand only the most promising candidates, leveraging the sorted property of each row. Instead of brute-forcing all combinations, we always expand the next smallest sum and use a visited set to avoid redundant work. This approach is both elegant and practical, making it a valuable technique for similar "k-th smallest combination" problems.