LeetCode 378: Kth Smallest Element in Sorted Matrix

Posted 2020-10-10 keepshuatishuati

1. Use priority queue. Need to check whether one element has been double counted:

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        if (matrix.length == 0 || matrix[0].length == 0) return 0;
        PriorityQueue<int[]> queue = new PriorityQueue<>((n1, n2) -> n1[0] - n2[0]);
        queue.offer(new int[]{matrix[0][0], 0, 0});
        int i = 1;
        int[] current = new int[3];
        Set<Integer> visited = new HashSet<>();
        while (!queue.isEmpty() && i++ <= k) {
            current = queue.poll();
            if (current[1] < matrix.length - 1 && !visited.contains((current[1] + 1)*matrix[0].length + current[2])) {
                queue.offer(new int[] {matrix[current[1] + 1][current[2]], current[1] + 1, current[2]});
                visited.add((current[1] + 1)*matrix[0].length + current[2]);
            }
            if (current[2] < matrix[0].length - 1 && !visited.contains(current[1]*matrix[0].length + current[2] + 1)) {
                queue.offer(new int[] {matrix[current[1]][current[2] + 1], current[1], current[2] + 1});
                visited.add(current[1] * matrix[0].length + current[2] + 1);
            }
        }
        return current[0];
    }
}

 

 

2 Binary search: For this kind of matrix, binary search should work as counting how many number satify the condition for each column and row.

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        if (matrix.length == 0 || matrix[0].length == 0) return 0;
        int start = matrix[0][0], end = matrix[matrix.length - 1][matrix[0].length - 1];
        while (start < end) {
            int mid = start + (end - start) / 2;
            int count = 0, j = matrix[0].length - 1;
            for (int i = 0; i < matrix.length; i++) {
                while (j >= 0 && matrix[i][j] > mid) j--;
                count += j + 1;
            }
            
            if (count < k) start = mid + 1;
            else end = mid;
        }
        return start;
    }
}