题目描述
链接:https://www.nowcoder.com/questionTerminal/7201cacf73e7495aa5f88b223bbbf6d1
给定两个有序数组arr1和arr2,再给定一个整数k,返回来自arr1和arr2的两个数相加和最大的前k个,两个数必须分别来自两个数组
要求:
1. 按照降序输出
2. 时间复杂度为O(klogk)
题目分析
这道题的正确解法为使用优先队列.
我们已知,两个数组是有序的, 然而要求递减的两个数组成的最大和, 那么必然有一定的规律性.
- 右下角的数必定是最大的
- 我们每次需要判断的是: 右下角的上方和左方(固定最大行和固定最大列, 哪个更大)
- 利用优先队列做BFS搜索即可, 如果找到了K个数, 直接return. 复杂度为KlogK.
代码 (注意边界)
import java.util.*;
public class Code04_TopKSumCrossTwoArrays {
public static class Node {
int x;
int y;
int sum;
public Node(int x, int y, int sum) {
this.x = x;
this.y = y;
this.sum = sum;
}
}
public static class MaxHeapComp implements Comparator<Node> {
@Override
public int compare(Node o1, Node o2) {
return o2.sum - o1.sum;
}
}
public static int[] topKSum(int[] arr1, int[] arr2, int topK) {
if (arr1 == null || arr2 == null || topK < 1) return null;
topK = Math.min(topK, arr1.length * arr2.length);
int[] res = new int[topK];
int resIndex = 0;
// 建立大根堆
PriorityQueue<Node> maxHeap = new PriorityQueue<>(new MaxHeapComp());
boolean[][] visit = new boolean[arr1.length][arr2.length];
int i1 = arr1.length - 1;
int i2 = arr2.length - 1;
// 初始化队列
maxHeap.add(new Node(i1, i2, arr1[i1] + arr2[i2]));
visit[i1][i2] = true;
while (resIndex != topK) {
Node curNode = maxHeap.poll();
if (curNode == null) break;
res[resIndex] = curNode.sum;
i1 = curNode.x;
i2 = curNode.y;
if (i1 - 1 >= 0 && !visit[i1-1][i2]) {
visit[i1-1][i2] = true;
maxHeap.add(new Node(i1 - 1, i2, arr1[i1-1] + arr2[i2]));
}
if (i2 - 1 >= 0 && !visit[i1][i2-1]) {
visit[i1][i2-1] = true;
maxHeap.add(new Node(i1, i2-1, arr1[i1] + arr2[i2-1]));
}
resIndex++;
}
return res;
}
public static void main(String[] args) {
int n, k;
Scanner sc = new Scanner(System.in);
n = sc.nextInt();
k = sc.nextInt();
int[] arr1 = new int[n];
int[] arr2 = new int[n];
for (int i = 0; i < n; i++) {
arr1[i] = sc.nextInt();
}
for (int i = 0; i < n; i++) {
arr2[i] = sc.nextInt();
}
int[] res = topKSum(arr1, arr2, k);
for (int i = 0; i < res.length; i++) {
if (i == 0) System.out.print(res[i]);
else System.out.print(" " + res[i]);
}
System.out.println();
}
}