Java PriorityQueue Solution


  • 3
    H

    32ms, O(k*log(k) time complexity, used inner class to simplify things and created a custom comparator for it.

    Explanation: Basically, we create a priorityqueue which will sort the inserted elements based on lowest matrix[][] value. At first, we insert the smallest value which is at matrix[0][0]. To iteratively find the next smallest element, we need to take out the smallest value from the priorityqueue, and add the value toward its right / bottom(if we are at the first row) into the priorityqueue (as long as they are within range). The next smallest value must be either already inside the priorityqueue, or be to the right/bottom of the current smallest value. Either way, the priorityqueue will make the correct ordering so that the next smallest value is ready to be removed. We do this loop k-1 times, and remove the kth smallest from the priorityqueue in the end.

    public class Solution {
        public class Point {
            int val;
            int x;
            int y;
            public Point(int val, int x, int y) {
                this.val = val;
                this.x = x;
                this.y = y;
            }
        }
        public int kthSmallest(int[][] matrix, int k) {
            if(matrix.length == 0) return 0;
            PriorityQueue<Point> pq = new PriorityQueue<Point>(new Comparator<Point>(){
               @Override
               public int compare(Point a, Point b) {
                   return a.val - b.val;
               }
            });
            pq.offer(new Point(matrix[0][0], 0, 0));
            for(int i = 1; i < k; i++) {
                Point p = pq.poll();
                if((p.x+1) < matrix.length) {
                    pq.offer(new Point(matrix[p.x+1][p.y], p.x+1, p.y));
                }
                if(p.x == 0 && (p.y+1) < matrix.length) {
                    pq.offer(new Point(matrix[p.x][p.y+1], p.x, p.y+1));
                }
            }
            return pq.poll().val;
        }
    }
    

  • 0
    A

    Could you also post some kind of explanation to make the code easier to understand? Thanks


  • 0
    H

    @aoben10 done, let me know if anything is unclear!


  • 2

    Fantastic solution, thank you so much!

    Here is my slightly cleaner version for those who are still a little confused:

    public class Solution {
        private class Value implements Comparable<Value> {
            int val;
            int row;
            int col;
            public Value(int val, int row, int col) {
                this.val = val;
                this.row = row;
                this.col = col;
            }
            public int compareTo(Value other) {
                return this.val - other.val;
            }
        }
        public int kthSmallest(int[][] matrix, int k) {
            PriorityQueue<Value> minHeap = new PriorityQueue<Value>();
            minHeap.add(new Value(matrix[0][0], 0, 0));
            for(int x = 1; x < k; x++) {
                Value val = minHeap.poll();
                if(val.row + 1 < matrix.length) {
                    minHeap.add(new Value(matrix[val.row + 1][val.col], val.row + 1, val.col));
                }
                // avoid duplicates
                if(val.row == 0 && val.col + 1 < matrix[0].length) {
                    minHeap.add(new Value(matrix[0][val.col + 1], 0, val.col + 1));
                }
            }
            return minHeap.peek().val;
        }
    }
    

  • 0
    H

    @DeusVult Yours is much cleaner and shows the intent better, thanks!


Log in to reply
 

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