Find the Kth Smallest Sum of a Matrix With Sorted Rows LeetCode solution in Java

1
class Solution {
2
  public int kthSmallest(int[][] mat, int k) {
3
    int[] row = mat[0];
4

5
    for (int i = 1; i < mat.length; i++) {
6
      row = findKthSmallest(row, mat[i], k);
7
    }
8

9
    return row[k - 1];
10
  }
11

12
  private int[] findKthSmallest(int[] num1, int[] num2, int k) {
13
    List<Integer> list = new ArrayList<>();
14
    PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (a[0] + a[1]) - (b[0] + b[1]));
15

16
    for (int i = 0; i < num1.length && i < k; i++) {
17
      minHeap.offer(new int[] {num1[i], num2[0], 0});
18
    }
19

20
    for (int i = 0; i < k && !minHeap.isEmpty(); i++) {
21
      int[] candidate = minHeap.poll();
22
      list.add(candidate[0] + candidate[1]); // SUM;
23
      int num2Idx = candidate[2];
24

25
      if (num2Idx < num2.length - 1) {
26
        minHeap.offer(new int[] {candidate[0], num2[num2Idx + 1], num2Idx + 1});
27
      }
28
    }
29

30
    return list.stream().mapToInt(i -> i).toArray();
31
  }
32
}