2015年6月10日星期三

lintcode @401 Kth Smallest Number in sorted matrix Java

Use a  minheap to keep a size of min( row, col , k) heap. row search or column search  k times.

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */
    public int kthSmallest(int[][] matrix, int k) {
 
        PriorityQueue<cell> pq= new PriorityQueue< cell>();
        cell res=new cell(0, 0, matrix[0][0]);

        if(matrix.length > matrix[0].length) {
            //serach by col
            for(int i=0;i< Math.min (matrix[0].length, k); i++)
             pq.offer(new cell(0, i, matrix[0][i]));
     
            while(k-- >0){
                res= pq. poll();
         
            if (res.x+1<matrix.length)
            pq.offer(new cell(res.x+1, res.y, matrix[res.x+1][res.y]));
            }
        }
        else {
            // height is smaller
            for(int i=0;i< Math.min(matrix.length, k); i++)
                pq.offer(new cell(i, 0, matrix[i][0]));
           
             while(k -- >=0){
                 res= pq. poll();
         
                if (res.y+1<matrix[0].length)
                pq.offer(new cell(res.x, res.y+1, matrix[res.x][res.y+1]));
            }
        }
        return res.d;
    }
   
    private class cell implements Comparable<cell>{
        public final int x;
        public final int y;
        public final  int d;
       
        public cell(int x, int y, int d){
            this.x=x;
            this.y=y;
            this.d=d;
        }
       
        public int compareTo(cell that){
            return d-that.d;
        }
       
    }
}

没有评论:

发表评论