January 06, 2022

378. Kth Smallest Element in a Sorted Matrix

378. Kth Smallest Element in a Sorted Matrix

Solution1. MaxHeap

Not optimized solution since it does not make use of the fact that rows and cols are sorted.

Time complexity: traverse all elements from array -> O(N^2*logK) Space complexity: O(K)

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        Queue<Integer> maxHeap = new PriorityQueue<>(k, (o1, o2) -> o2 - o1);
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[0].length; j++) {
                int num = matrix[i][j];
                maxHeap.add(num);
                if (maxHeap.size() > k) {
                    maxHeap.poll();
                }
            }
        }
        return maxHeap.peek();
    }
}

Solution2. minHeap


class Solution {
    static class HeapNode {
        int row;
        int col;
        int value;
        HeapNode(int row, int col, int value) {
            this.row = row;
            this.col = col;
            this.value = value;
        } 
    }
    
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        
        // Best First Search, need a minheap on the value of each node.
        PriorityQueue<HeapNode> minHeap = new PriorityQueue<>(k, (n1, n2) -> n1.value - n2.value);
        boolean[][] visited = new boolean[n][n];
        minHeap.offer(new HeapNode(0, 0, matrix[0][0]));
        visited[0][0] = true;
        
        // iterate k - 1 rounds, and Best First Search the smallest k-1 node.
        for (int i = 0; i < k - 1; i ++) {
            HeapNode cur = minHeap.poll();
            if (cur.row + 1 < n && !visited[cur.row + 1][cur.col]) {
                minHeap.offer(new HeapNode(cur.row + 1, cur.col, matrix[cur.row + 1][cur.col]));
                visited[cur.row + 1][cur.col] = true;
            } 
            if (cur.col + 1 < n && !visited[cur.row][cur.col + 1]) {
                minHeap.offer(new HeapNode(cur.row, cur.col + 1, matrix[cur.row][cur.col + 1]));
                visited[cur.row][cur.col + 1] = true;
            }
        }
        return minHeap.peek().value;
    }
}
comments powered by Disqus