Kth Elements in some Sorted Arrays
简化版:https://www.geeksforgeeks.org/k-th-element-two-sorted-arrays/
用三元组记录坐标和value,很巧妙,不然要用class。
public int mergeSort(int[][] nums, int k) {
int rowNum = nums.length;
PriorityQueue<int[]> pq = new PriorityQueue<>(rowNum, new Comparator<int[]>() {
public int compare(int[] i1, int[] i2) {
return i1[2] - i2[2];
}
});
// 把每一行的第一个元素,加入到优先队列中,count记录加入了多少个元素
int count = 0;
// sum用来数二维数组中一共有多少个元素
int sum = 0;
for (int i = 0; i < nums.length; i++) {
if (nums[i].length > 0) {
int[] arr = new int[]{i, 0, nums[i][0]};
pq.offer(arr);
count++;
sum += nums[i].length;
}
}
// 如果元素个数小于k,说明第k小的,不在数组中
if (sum < k) {
return Integer.MIN_VALUE;
}
// 跳过第一个if,说明,元素一定在数组中
// 如果加入的个数大于k,说明需要的元素已经在队列中了,只需取出来即可
if (count >= k) {
for (int i = 1; i < k; i++) {
pq.poll();
}
return pq.poll()[2];
}
// 跳过第二个if,说明需要的元素在数组中,但是还没有加入到优先队列中
// 此时,一定要从优先队列中出去的元素来进行记录,出去3个,说明找到了前3个最小的元素,count用来记录出去的个数
count = 0;
int[] head = null;
while (count != k) {
head = pq.poll();
count++;
int row = head[0];
int col = head[1];
if (col + 1 < nums[row].length) {
int[] next = new int[]{row, col + 1, nums[row][col + 1]};
pq.offer(next);
}
}
return head[2];
}