【牛客网】两个有序数组间相加和的Topk问题 [H]（堆）

两个有序数组间相加和的Topk问题_牛客题霸_牛客网 (nowcoder.com)

一、题目

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

按照降序输出。

要求：时间复杂度为O(k*log k)

示例 1：
输入：
5 4
1 2 3 4 5
3 5 7 9 11
输出：
16 15 14 14

二、代码

import java.util.Scanner;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.StreamTokenizer;
import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
// 注意类名必须为 Main, 不要有任何 package xxx 信息
public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StreamTokenizer in = new StreamTokenizer(br);
        PrintWriter out = new PrintWriter(new OutputStreamWriter(System.out));
        while (in.nextToken() != StreamTokenizer.TT_EOF) {
            int N = (int) in.nval;
            in.nextToken();
            int K = (int) in.nval;
            int[] arr1 = new int[N];
            int[] arr2 = new int[N];
            for (int i = 0; i < N; i++) {
                in.nextToken();
                arr1[i] = (int) in.nval;
            }
            for (int i = 0; i < N; i++) {
                in.nextToken();
                arr2[i] = (int) in.nval;
            }
            int[] topK = topKSum(arr1, arr2, K);
            for (int i = 0; i < K; i++) {
                out.print(topK[i] + " ");
            }
            out.println();
            out.flush();
        }
    }

    public static class Node {
        // 当前遍历到的arr1的下标，也就是矩阵的第一个下标
        public int i;
        // 当前遍历到的arr2的下标，也就是矩阵的第二个下标
        public int j;
        // 两个数组成的加和
        public int sum;

        public Node(int i, int j, int s) {
            this.i = i;
            this.j = j;
            this.sum = s;
        }
    }

    // 构造大根堆比较器，加和大的在堆顶
    public static class MaxHeapComp implements Comparator<Node> {
        @Override
        public int compare(Node o1, Node o2) {
            return o2.sum - o1.sum;
        }
    }

    // 返回加和前k大的结果
    public static int[] topKSum(int[] arr1, int[] arr2, int k) {
        // 矩阵右下角就是加和的最大值，因为两个数组是有序的，最后一个数一定是最大的
        int i = arr1.length - 1;
        int j = arr2.length - 1;

        // 标记当前矩阵位置是否已经加入过大根堆
        // 这里使用矩阵的二位坐标转一维坐标来标记的
        HashSet<Integer> set = new HashSet<>();
        // 创建大根堆
        PriorityQueue<Node> maxHead = new PriorityQueue<>(new MaxHeapComp());
        // 先将矩阵右下角的数加入大根堆
        maxHead.add(new Node(i, j, arr1[i] + arr2[j]));
        // 标记这个位置已经加入过大根堆了，二维坐标转一维坐标
        set.add(i * arr2.length + j);
        // 记录答案
        int[] ans = new int[k];
        int index = 0;
        // 如果k大于所有数的总数量，那么topK就设置为所有数的总数量
        int topK = Math.min(k, arr1.length * arr2.length);

        // 收集完topK就结束循环
        while (index != topK) {
            // 弹出堆顶，记录一个答案
            Node node = maxHead.poll();
            ans[index++] = node.sum;

            // 将弹出堆顶的上面和左边的位置加入到大根堆中，需要判断上面和下面是否还有数，并且这个数还不能加入过大根堆
            if (node.i > 0 && !set.contains((node.i - 1) * arr2.length + node.j)) {
                // 将上面的数加入大根堆
                maxHead.add(new Node(node.i - 1, node.j, arr1[node.i - 1] + arr2[node.j]));
                // 标记已经加入过堆了
                set.add((node.i - 1) * arr2.length + node.j);
            }

            if (node.j > 0 && !set.contains(node.i * arr2.length + (node.j - 1))) {
                // 将左面的数加入大根堆
                maxHead.add(new Node(node.i, node.j - 1, arr1[node.i] + arr2[node.j - 1]));
                 // 标记已经加入过堆了
                set.add(node.i * arr2.length + (node.j - 1));
            }
        }
        // 返回答案
        return ans;
    }

   
}

三、解题思路 

大根堆，想象出一个矩阵，但是代码并不会真的创建一个矩阵，还是在两个数组上进行遍历。arr1做行对应，arr2做列对应。

整个题就要利用两个数组的有序性。先去右下角的最大值，将其加入到大根堆中。此时大根堆中只有一个数，堆顶的就是目前找到的最大的数。

将堆顶数据弹出，记录下一个答案。

然后将该堆顶在矩阵中地所在位置的上面和左边的两组加和加入到大根堆中。这两个数中一定有仅次于刚才从堆中弹出那个数的，所以将他俩加入大根堆，此时的堆顶就是我们要的另一个答案。

整个过程就是在利用原数组的单调性，每次都一点点的减少横纵下标，而且会挑着比较大的那个位置将其上面和左边的加入到堆中，这样就能保证不会遗漏更大的数了，每次都是从最大的那个加和组合基础上来减少参与加和的两个数的大小。

注意：防止同一个位置进入堆。要保证进大根堆不要重复进。是有可能出现这种情况的，所以要记录上哪些数已经进入过大根堆了。