Java heap solution, time complexity klog(k)


  • 4
    W

    In heap, element is an array
    int[0] represents matrix[i][j] value
    int[1] represents i
    int[2] represents j
    Whenever an element is poll() from heap, push the element below it to heap, only push the element right to it to heap when it's in first row. So we can avoid duplicates.

    public int kthSmallest(int[][] matrix, int k) {
            PriorityQueue<int[]> heap = new PriorityQueue<int[]>(10000,new Comparator<int[]>(){
                @Override
                public int compare(int[] a,int[] b){
                    return a[0]-b[0];
                }
            });
            if(matrix.length==0||k==0){
                return -1;
            }
            heap.offer(new int[]{matrix[0][0],0,0});
            int[] peak = new int[3];
            while(k-->0){
                peak = heap.poll();
                if(peak[1]+1<matrix.length){
                    heap.offer(new int[]{matrix[peak[1]+1][peak[2]],peak[1]+1,peak[2]});
                }
                if(peak[1]==0&&peak[2]+1<matrix[0].length){
                    heap.offer(new int[]{matrix[peak[1]][peak[2]+1],peak[1],peak[2]+1});
                }
            }
            return peak[0];
        }
    

    Still thinking about binary search solution.


  • 0
    U
    This post is deleted!

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.