LeetCode 378. Kth Smallest Element in a Sorted Matrix

Question

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(n^2).

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].length
• 1 <= n <= 300
• -109 <= matrix[i][j] <= 109
• All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
• 1 <= k <= n2

Source: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/

Solution

Generally, there are three ways to find the kth smallest/largest element: 1. Max Heap 2. Min Heap 3. Quick Sort. We will skip quicksort for now in this blog post.

If using max heap, we should instantiate a max heap with size of k, and fill the heap with the first k elements of the matrix (it is the initilization process). Then we traverse the rest of elements in the matrix. If an element e is greater than the heap top, drop it; otherwise, pop the heap top and push element e. After this process, elements in the max heap are k-smallest elements. And the heap top is the result. However, because we need to traverse the whole matrix (can be optimized but still on the order of O(n^2)) and do not utilize the order feature of given matrix, the time complexity will be O(n^2*log(k)).

In the other solution, we treat a sorted matrix as n sorted list. So it can be regarded as an extension of finding kth smallest element in two sorted list. The difference is that we do not use n variables to store indexes on list and compare them, but use a min heap to do this job.

The main difference between max heap solution and min heap solution:

We use the max heap to find k-smallest elements by traverse the matrix and comparing elements with the heap top. For min heap, we get k-smallest elements by other means and use min heap to find largest one of them (the kth smallest element).

class element {
    // matrix[x][y]
    int val;
    int x;
    int y;

    public element(int _val, int _x, int _y) {
        this.val = _val;
        this.x = _x;
        this.y = _y;
    }
}

// min heap, Time Complexity O(k*log(k))
public int kthSmallest2(int[][] matrix, int k) {
    int n = matrix.length;
    int numSortedRow = Math.min(n, k);
    PriorityQueue<element> minHeap = new PriorityQueue<>(numSortedRow, new Comparator<element>() {
        @Override
        public int compare(element o1, element o2) {
            return o1.val - o2.val;
        }
    });
    for (int i = 0; i < numSortedRow; i++) {
        minHeap.add(new element(matrix[i][0], i, 0));
    }

    int result = 0;
    // treat a sorted matrix as n sorted lists
    for (int count = 0; count < k; count++) {
        element e = minHeap.poll();
        result = e.val;
        if (e.y < n - 1) {
            minHeap.add(new element(matrix[e.x][e.y + 1], e.x, e.y + 1));
        }
    }
    return result;
}