Java heap solution with comments


  • 0
    A
    class Solution {
        
        //each node has a value (in the matrix) a row and an index in that row (column)
        class Node {
            int value;
            int row;
            int idx;
            public Node(int v, int r, int i) {
                value = v;
                row = r;
                idx = i;
            }
        }
        
        public int kthSmallest(int[][] matrix, int k) {
            int count = 0;
            PriorityQueue<Node> queue = new PriorityQueue<Node>(matrix.length, new Comparator<Node>() {
                //compare nodes based off value in the matrix
                public int compare(Node n1, Node n2) {
                    return n1.value - n2.value;
                }
            });
            
            //add the first column elements into the heap
            for(int x = 0; x < matrix.length; x++) {
                queue.add(new Node(matrix[x][0], x, 0));
            }
            
            while(count != k - 1) {
                Node n = queue.poll();
                //make sure the idx does not go past the length of the array in the matrix
                if(n.idx + 1 < matrix[n.idx].length) {
                    queue.add(new Node(matrix[n.row][n.idx + 1], n.row, n.idx + 1));
                }
                count++;
            }
            
            return queue.poll().value;
        }
    }
    

Log in to reply
 

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