两个有序数组间相加和的TopK问题 - 优先队列

题目描述

链接：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();
    }
}