一.题目描述
给定两个有序数组arr1和arr2,再给定一个整数k,返回来自arr1和arr2的两个数相加和最大的前k个,两个数必须分别来自两个数组
按照降序输出
二.代码
import java.util.*;
public class Main {
static class Node {
public int index1; //arr1中的位置
public int index2; //arr2中的位置
public int sum; //arr1[index1]+arr2[index2]
public Node(int i1, int i2, int s) {
index1 = i1;
index2 = i2;
sum = s;
}
}
public static int[] topKSum(Integer[] arr1, Integer[] arr2, int topK) {
topK = Math.min(topK, arr1.length * arr2.length);
int[] res = new int[topK];
int resIndex = 0;
PriorityQueue<Node> maxHeap = new PriorityQueue<>((N1, N2) -> N2.sum - N1.sum);
HashSet<String> positionSet = new HashSet<>();
int i1 = arr1.length - 1;
int i2 = arr2.length - 1;
maxHeap.add(new Node(i1, i2, arr1[i1] + arr2[i2]));
positionSet.add(i1 + "_" + i2);
while (resIndex != topK) {
Node curNode = maxHeap.poll();
res[resIndex++] = curNode.sum;
i1 = curNode.index1;
i2 = curNode.index2;
if (i1 - 1 >= 0 && !positionSet.contains(i1 - 1 + "_" + i2)) {
positionSet.add(i1 - 1 + "_" + i2);
maxHeap.add(new Node(i1 - 1, i2, arr1[i1 - 1] + arr2[i2]));
}
if (i2 - 1 >= 0 && !positionSet.contains(i1 + "_" + (i2 - 1))) {
positionSet.add(i1 + "_" + (i2 - 1));
maxHeap.add(new Node(i1, i2 - 1, arr1[i1] + arr2[i2 - 1]));
}
}
return res;
}
public static void main(String[] args) {
Scanner in = new Scanner(System.in);
int n = in.nextInt();
int k = in.nextInt();
Integer[] arr1 = new Integer[n];
Integer[] arr2 = new Integer[n];
for (int i = 0; i < n; i++) {
arr1[i] = in.nextInt();
}
for (int i = 0; i < n; i++) {
arr2[i] = in.nextInt();
}
int[] res = topKSum(arr1, arr2, k);
for (int re : res) {
System.out.print(re + " ");
}
}
}