Problem Statement:
Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.
Note that it is the kth smallest element in the sorted order, not the kth distinct element.
You must find a solution with a memory complexity better than O(n2).
Examples:
Example 1:
Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output:13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13
Example 2:
Input:matrix = [[-5]], k = 1
Output:-5
Constraints:
n == matrix.length == matrix[i].length1 <= n <= 300-109 <= matrix[i][j] <= 109.- It is guaranteed that the answer is unique.
- All the rows and columns of
matrixare guaranteed to be sorted in non-decreasing order. 1 <= k <= n2
Approach
- Treat each row of the sorted matrix like a
sortedlist. - Use a
min-heaptostore the smallest availableelement from each row. - Initially, push the first element of each of the first
min(n, k)rows into the heap. - Repeatedly pop the smallest element from the heap and, if possible, push the next element from the same row.
- After popping
k-1 times, the top of the heap is the k-th smallest element.
Time Complexity:
Time Complexity = 0(min(n,k) log(min(n,k))) = O(n log n)
Space Complexity:
Space Complexity = 0(Min(n,k))
Dry Run
Input: stones = [2, 7, 4, 1, 8, 1]
State Transitions:
Start Function lastStoneWeight([2, 7, 4, 1, 8, 1]) pq = new MaxPriorityQueue() -- Enqueue all stones -- i = 0 → stones[0] = 2 → pq.enqueue(2) → pq = [2] i = 1 → stones[1] = 7 → pq.enqueue(7) → pq = [7, 2] i = 2 → stones[2] = 4 → pq.enqueue(4) → pq = [7, 2, 4] i = 3 → stones[3] = 1 → pq.enqueue(1) → pq = [7, 2, 4, 1] i = 4 → stones[4] = 8 → pq.enqueue(8) → pq = [8, 7, 4, 2, 1] i = 5 → stones[5] = 1 → pq.enqueue(1) → pq = [8, 7, 4, 2, 1, 1] -- Enter while loop (pq.size() > 1) -- Iteration 1: → y = pq.dequeue() -> 8 → x = pq.dequeue() -> 7 → y - x = 1 ( > 0 ) → pq.enqueue(1) pq after enqueue -> [4, 2, 1, 1, 1] Iteration 2: → y = pq.dequeue() -> 4 → x = pq.dequeue() -> 2 → y - x = 2 ( > 0 ) → pq.enqueue(2) pq after enqueue -> [2, 1, 1, 1, 2] (equivalently [2, 2, 1, 1, 1]) Iteration 3: → y = pq.dequeue() -> 2 → x = pq.dequeue() -> 2 → y - x = 0 (== 0) → do NOT enqueue pq after removals -> [1, 1, 1] Iteration 4: → y = pq.dequeue() -> 1 → x = pq.dequeue() -> 1 → y - x = 0 (== 0) → do NOT enqueue pq after removals -> [1] -- Exit while (pq.size() == 1) -- Return pq.dequeue() || 0 -> returns 1
Final Output: 1
Explanation: repeatedly smash the two largest stones; after all smashes one stone of weight 1 remains, so the function returns 1.
var kthSmallest = function(matrix, k) {
let n = matrix[0].length;
let heap = new MinPriorityQueue(x => x.val);
for(let i = 0; i < Math.min(n,k); i++) {
heap.push({val: matrix[i][0], row: i, col: 0});
}
for(let count=0; count < k-1; count++) {
let {val, row, col} = heap.pop();
if(col+1 < n) {
heap.push({val: matrix [row][col+1], row: row, col: col+1});
}
}
return heap.pop().val;
};
import heapq
def kthSmallest(matrix, k):
n = len(matrix)
heap = []
for i in range(min(n, k)):
heapq.heappush(heap, (matrix[i][0], i, 0))
for _ in range(k - 1):
val, row, col = heapq.heappop(heap)
if col + 1 < n:
heapq.heappush(heap, (matrix[row][col + 1], row, col + 1))
return heapq.heappop(heap)[0]
class Solution {
static class Node {
int val, row, col;
Node(int v, int r, int c) {
val = v; row = r; col = c;
}
}
public int kthSmallest(int[][] matrix, int k) {
int n = matrix.length;
PriorityQueue pq = new PriorityQueue<>((a, b) -> a.val - b.val);
for (int i = 0; i < Math.min(n, k); i++) {
pq.offer(new Node(matrix[i][0], i, 0));
}
for (int count = 0; count < k - 1; count++) {
Node top = pq.poll();
if (top.col + 1 < n) {
pq.offer(new Node(matrix[top.row][top.col + 1], top.row, top.col + 1));
}
}
return pq.poll().val;
}
}
struct Node {
int val, row, col;
bool operator>(const Node &other) const {
return val > other.val;
}
};
int kthSmallest(vector>& matrix, int k) {
int n = matrix.size();
priority_queue, greater> minHeap;
for (int i = 0; i < min(n, k); i++) {
minHeap.push({matrix[i][0], i, 0});
}
for (int count = 0; count < k - 1; count++) {
Node top = minHeap.top();
minHeap.pop();
if (top.col + 1 < n) {
minHeap.push({matrix[top.row][top.col + 1], top.row, top.col + 1});
}
}
return minHeap.top().val;
}
typedef struct {
int val, row, col;
} Node;
void swap(Node* a, Node* b) {
Node temp = *a; *a = *b; *b = temp;
}
void heapify(Node heap[], int n, int i) {
int smallest = i, l = 2*i + 1, r = 2*i + 2;
if (l < n && heap[l].val < heap[smallest].val) smallest = l;
if (r < n && heap[r].val < heap[smallest].val) smallest = r;
if (smallest != i) {
swap(&heap[i], &heap[smallest]);
heapify(heap, n, smallest);
}
}
int kthSmallest(int** matrix, int n, int k) {
Node heap[k];
int size = 0;
for (int i = 0; i < (n < k ? n : k); i++) {
heap[size++] = (Node){matrix[i][0], i, 0};
}
for (int i = size / 2 - 1; i >= 0; i--) heapify(heap, size, i);
for (int count = 0; count < k - 1; count++) {
Node top = heap[0];
if (top.col + 1 < n) {
heap[0] = (Node){matrix[top.row][top.col + 1], top.row, top.col + 1};
} else {
heap[0] = heap[size - 1];
size--;
}
heapify(heap, size, 0);
}
return heap[0].val;
}
using System;
using System.Collections.Generic;
public class Solution {
public class Node {
public int val, row, col;
public Node(int v, int r, int c) { val = v; row = r; col = c; }
}
public int KthSmallest(int[][] matrix, int k) {
int n = matrix.Length;
var pq = new PriorityQueue();
for (int i = 0; i < Math.Min(n, k); i++) {
pq.Enqueue(new Node(matrix[i][0], i, 0), matrix[i][0]);
}
for (int count = 0; count < k - 1; count++) {
var top = pq.Dequeue();
if (top.col + 1 < n) {
pq.Enqueue(new Node(matrix[top.row][top.col + 1], top.row, top.col + 1),
matrix[top.row][top.col + 1]);
}
}
return pq.Dequeue().val;
}
}
