Kth Smallest Number in Sorted Matrix

ID: 401; medium

LintCode 炼码 - 更高效的学习体验！www.lintcode.com

Solution 1 (Java)

public class Solution {

    class Pair {
        int x, y, val;
        public Pair(int x, int y, int val) {
            this.x = x;
            this.y = y;
            this.val = val;
        }
    }
    /**
     * @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) {
        int[] dx = {0, 1};
        int[] dy = {1, 0};
        int n = matrix.length, m = matrix[0].length;
        boolean[][] visited = new boolean[n][m];
        Queue<Pair> minHeap = new PriorityQueue<>(new Comparator<Pair>() {
            public int compare(Pair p1, Pair p2) {
                return p1.val - p2.val;
            }
        });
        minHeap.offer(new Pair(0, 0, matrix[0][0]));

        for (int i = 0; i < k - 1; i++) {
            Pair curr = minHeap.poll();
            for (int j = 0; j < 2; j++) {
                int next_x = curr.x + dx[j];
                int next_y = curr.y + dy[j];
                if (next_x < n && next_y < m && !visited[next_x][next_y]) {
                    Pair next = new Pair(next_x, next_y, matrix[next_x][next_y]);
                    minHeap.offer(next);
                    visited[next_x][next_y] = true;
                }
            }
        }
        return minHeap.peek().val;
    }
}

Notes

• We use a min heap to keep track of the numbers. This solution is essentially a BFS. Starting from the top left number, which is the minimum number in the matrix, then we add the number on its right and the number below it.

• We poll the min heap k - 1 times and eventually we end up with the kth smallest number in the heap.

• It is also important to keep track of which numbers are already visited, just like what is down in graph traversal problems;

• Time complexity: O(klogn)

Solution 2 (Java)

public class Solution {

    class Result {
        boolean exists;
        int rank;
        public Result(boolean e, int r) {
            exists = e;
            rank = r;
        }
    }
    /**
     * @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) {
        int n = matrix.length, m = matrix[0].length;
        int left = matrix[0][0];
        int right = matrix[n - 1][m - 1];
        while (left <= right) {
            int mid = left + (right - left) / 2;
            Result res = check(mid, matrix);
            if (res.exists && res.rank == k) {
                return mid;
            } else if (res.rank < k) {
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return left;
    }

    private Result check(int value, int[][] matrix) {
        int n = matrix.length, m = matrix[0].length;
        int i = n - 1, j = 0;
        Result res = new Result(false, 0);
        while (i >= 0 && j < m) {
            if (matrix[i][j] == value)
                res.exists = true;
            if (matrix[i][j] <= value) {
                res.rank += i + 1;
                j += 1;
            } else {
                i -= 1;
            }
        }
        return res;
    }

    
}

Notes

• Using the sorted property, it is natural to think about binary search.