牛客网-两个有序数组间相加和的Topk问题

一.题目描述

给定两个有序数组arr1和arr2，再给定一个整数k，返回来自arr1和arr2的两个数相加和最大的前k个，两个数必须分别来自两个数组

按照降序输出

两个有序数组间相加和的Topk问题__牛客网

二.代码

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 + " ");
        }
    }
}