Kth Smallest Number in Sorted Matrix 401

Question

Find the kth smallest number in at row and column sorted matrix.

Example

Given k = 4 and a matrix:

[ [1 ,5 ,7],

[3 ,7 ,8],

[4 ,8 ,9], ]

return 5

Solution

用一个PriorityQueue(Heap)来实现。PriorityQueue堆顶保存的为当前元素中最小的元素:

1) 先将k个第一列元素加入PriorityQueue(若第一列元素个数小于k,加入全部第一列元素)。此时,堆顶元素为matrix[0][0],为当前最小元素,matrix[1][0]为第二小元素。

2) 将堆顶元素出队,加入其右边元素matrix[0][1]。若matrix[1][0]比matrix[0][1]小,则matrix[1][0]为堆顶元素,且为当前元素中最小元素(比它下面,右边,matrix[0][1]及其右边元素都小);反之,则matrix[0][1]为堆顶元素,且为当前元素中最小元素。因此,将出队元素(当前最小值)的右边元素入队已定能保证该性质。

3) 重复2)中过程,每次出队的都是当前元素中最小元素,出队k-1此之后的堆顶元素即为第k个元素。

代码如下:

public class Solution {
    /**
     * @param matrix: a matrix of integers
     * @param k: an integer
     * @return: the kth smallest number in the matrix
     */

    public class Point {
        public int x, y, val;
        public Point(int x, int y, int val) {
            this.x = x;
            this.y = y;
            this.val = val;
        }
    } 

    Comparator<Point> comp = new Comparator<Point>() {
        public int compare(Point left, Point right) {
            return left.val - right.val;
        }
    };

    public int kthSmallest(int[][] matrix, int k) {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0) {
            return 0;
        }
        if (k > matrix.length * matrix[0].length) {
            return 0;
        }
        return horizontal(matrix, k);
    }

    private int horizontal(int[][] matrix, int k) {
        Queue<Point> heap = new PriorityQueue<Point>(k, comp);
        for (int i = 0; i < Math.min(matrix.length, k); i++) {
            heap.offer(new Point(i, 0, matrix[i][0]));
        }
        for (int i = 0; i < k - 1; i++) {
            Point curr = heap.poll();
            if (curr.y + 1 < matrix[0].length) {
                heap.offer(new Point(curr.x, curr.y + 1, matrix[curr.x][curr.y + 1]));
            }
        }
        return heap.peek().val;
    }
 }

results matching ""

    No results matching ""