[LeetCode] Kth Smallest Element in a Sorted Matrix

Kth Smallest Element in a Sorted Matrix

Given a n x n matrix where each of the rows and columns are sorted in
ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the
kth distinct element.

BFS + Priority Queue

Time Complexity
O(klogk)

Space Complexity
每一轮最多出一个node,放两个Node, 进行了k轮
(-1 + 2)k O(k)
去重也用了数据结构, 如果是用了一个boolean matrix O(m*n)
使用hashset: HashSet 的size最多是把所有visited过的点丢进Hashset里面 + 加上在Heap中的
点
O(2k)

思路

根据给的matrix的性质，可以知道最小的element一定是在左上角的，接下来小的值是在这个值的右边还是在这个值的下面没有办法确定，这里维护的是一个size为k的minHeap，通过传统的BFS做法，不停地把当前值的下面的和左边的传入minHeap，因为minHeap会自动的把最小值放到最上面，所以poll k次就可以知道第k小的。这里用到一个visited的hashset是为了标记visited过的点，也可以设置一个booleanmatrix.length.length] visited，用Hashset可以节约一点空间。

代码

class Element{
    //fields
    private int val;
    private int x;  //row index of the current value in matrix
    private int y;  //col index of the current value in matrix
    
    //constructor
    public Element(int val, int x, int y){
        this.val = val;
        this.x = x;
        this.y = y;
    }
}
public int kthSmallest(int[][] matrix, int k) {
    //corner case
    if(matrix == null || matrix.length == 0 || matrix[0] == null || matrix[0].length == 0){
        return -1;
    }
    int rows = matrix.length, cols = matrix[0].length;
    //use min heap, compare the value of the element in matrix
    Queue<Element> minHeap = new PriorityQueue<Element>(k, new Comparator<Element>(){
        @Override
        public int compare(Element e1, Element e2){
            return e1.val - e2.val;
        }
    });
    //use a hashset to check if the element has been visited
    Set<Integer> visited = new HashSet<Integer>();
    //offer the smallest element in the matrix to the minHeap first
    minHeap.offer(new Element(matrix[0][0], 0, 0));
    //put this element as visited
    visited.add(0);
    
    //we need to poll k times to find the kth smallest
    while(k > 1){
        //expand the priority queue
        Element cur = minHeap.poll();
        k--;
        //expand from the row
        //check if the index is valid, if it is in the matrix and if it is visited
        if(cur.y + 1 >= 0 && cur.y + 1 < cols && !visited.contains(cur.x * cols + cur.y + 1)){
            minHeap.offer(new Element(matrix[cur.x][cur.y + 1], cur.x, cur.y + 1));
            visited.add(cur.x * cols + cur.y + 1);
        }
        //expand from the col
        //do the same check
        if(cur.x + 1 >= 0 && cur.x + 1 < rows && !visited.contains((cur.x + 1) * cols + cur.y)){
            minHeap.offer(new Element(matrix[cur.x + 1][cur.y], cur.x + 1, cur.y));
            visited.add((cur.x + 1) * cols + cur.y);
        }
    }
    return minHeap.poll().val;
}