Kth Smallest Element in a Sorted Matrix LeetCode solution in Java

1
/**
2
 * Using PriorityQueue
3
 *
4
 * <p>Time Complexity: O(min(N,K)*log(min(N,K))) -> To add initial min(N,K) elements, as we are
5
 * adding the elements individually. If we were adding all elements in one go, then the complexity
6
 * would be O(min(N,K)) Refer: https://stackoverflow.com/a/34697891 O(2*(K-1)*log(min(N,K)) -> To
7
 * poll K-1 elements and add next K-1 elements. Total Time Complexity: O((min(N,K) + 2*(K-1)) *
8
 * log(min(N,K)) = O(K * log(min(N,K))
9
 *
10
 * <p>Space Complexity: O(min(N, K))
11
 *
12
 * <p>N = Length of one side of the matrix. K = input value k.
13
 */
14
class Solution {
15
  public int kthSmallest(int[][] matrix, int k) {
16
    if (matrix == null || k <= 0) {
17
      throw new IllegalArgumentException("Input is invalid");
18
    }
19

20
    int n = matrix.length;
21
    if (k > n * n) {
22
      throw new NoSuchElementException("k is greater than number of elements in matrix");
23
    }
24
    if (k == 1) {
25
      return matrix[0][0];
26
    }
27
    if (k == n * n) {
28
      return matrix[n - 1][n - 1];
29
    }
30

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

34
    for (int i = 0; i < Math.min(n, k); i++) {
35
      queue.offer(new int[] {i, 0});
36
    }
37
    while (k > 1) {
38
      int[] cur = queue.poll();
39
      if (cur[1] < n - 1) {
40
        cur[1]++;
41
        queue.offer(cur);
42
      }
43
      k--;
44
    }
45

46
    return matrix[queue.peek()[0]][queue.peek()[1]];
47
  }
48
}