378. Kth Smallest Element in a Sorted Matrix

// LeetCode URL: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/
import java.util.PriorityQueue;

/**
 * Using PriorityQueue
 *
 * Time Complexity: O(A log A + K log K)
 *
 * Space Complexity: O(A)
 *
 * A = min(N, K). N = size of matrix. K = input value k.
 */
class Solution1 {
    public int kthSmallest(int[][] matrix, int k) throws IllegalArgumentException {
        if (matrix == null || k <= 0 || matrix.length == 0 || matrix[0].length == 0) {
            throw new IllegalArgumentException("Illegal Input");
        }

        int n = matrix.length;

        if (k > n * n) {
            throw new IllegalArgumentException("k is greater than the number of elements in the matrix.");
        }
        if (k == 1) {
            return matrix[0][0];
        }
        if (k == n * n) {
            return matrix[n - 1][n - 1];
        }

        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> (matrix[a[0]][a[1]] - matrix[b[0]][b[1]]));

        for (int i = 0; i < Math.min(n, k); i++) {
            queue.offer(new int[] { i, 0 });
        }
        for (int i = 0; i < k - 1; i++) {
            int[] cell = queue.poll();
            if (cell[1] == n - 1) {
                continue;
            }
            queue.offer(new int[] { cell[0], cell[1] + 1 });
        }

        int[] res = queue.poll();
        return matrix[res[0]][res[1]];
    }
}

/**
 * Using Binary Search
 *
 * Refer: 1)
 * https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85173/Share-my-thoughts-and-Clean-Java-Code
 * 2)
 * https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85173/Share-my-thoughts-and-Clean-Java-Code/122873
 *
 * Time Complexity: O(N^2 * log(N^2)) = O(N^2 * log N)
 *
 * Space Complexity: O(1)
 *
 * N = size of matrix.
 */
class Solution2 {
    public int kthSmallest(int[][] matrix, int k) throws IllegalArgumentException {
        if (matrix == null || matrix.length == 0 || matrix[0].length == 0 || k <= 0) {
            throw new IllegalArgumentException("Illegal Input");
        }

        int n = matrix.length;

        if (k > n * n) {
            throw new IllegalArgumentException("k is greater than the number of elements in the matrix.");
        }
        if (k == 1) {
            return matrix[0][0];
        }
        if (k == n * n) {
            return matrix[n - 1][n - 1];
        }

        int low = matrix[0][0];
        int high = matrix[n - 1][n - 1];

        while (low < high) {
            int mid = low + (high - low) / 2;
            int count = 0;
            int j = n - 1;
            for (int i = 0; i < n; i++) {
                while (j >= 0 && matrix[i][j] > mid) {
                    j--;
                }
                count += (j + 1);
            }
            if (count < k) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }

        return low;
    }
}