083 两个有序数组间相加和的Topk问题
Table of Contents

ac

牛客网反应出来的问题注意点,也可以作为理解题目问答

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;


class DataNode{
    int i;
    int j;

    public DataNode(int i, int j) {
        this.i = i;
        this.j = j;
    }
}

public class Main {
    public static Scanner sc = new Scanner(System.in);

    public static BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));

    public static void main(String[] args) throws Exception {
        int n = sc.nextInt();
        int k = sc.nextInt();
        long[] a1 = new long[n];
        long[] a2 = new long[n];
        for(int i=0;i<n;i++){
            a1[i] = sc.nextLong();
        }
        for(int i=0;i<n;i++){
            a2[i] = sc.nextLong();
        }
        Arrays.sort(a1);
        Arrays.sort(a2);
        // 大根堆
        PriorityQueue<DataNode> priorityQueue = new PriorityQueue<>((a,b)->{
            long val1 = a1[a.i] + a2[a.j];
            long val2 = a1[b.i] + a2[b.j];
            if(val1 > val2){
                return -1;
            }
            if(val1 < val2){
                return 1;
            }
            return 0;
        });

//        boolean[][] posOccur = new boolean[n][n]; // n*n可能越界
//        posOccur[n-1][n-1] = true; // n * n 内存太多了,内存超限制

        Map<Integer,Map<Integer,Boolean>> mp = new HashMap<>();

        StringBuilder sb = new StringBuilder();
        priorityQueue.offer(new DataNode(n-1,n-1));
        int popNum = 0;
        while (popNum < k){
            if(!priorityQueue.isEmpty()){
                DataNode dataNode = priorityQueue.poll();
                long val = a1[dataNode.i] + a2[dataNode.j];
                if(popNum == 0){
//                    sb.append(val);
                    popNum ++;
                }else {
//                    sb.append(" " + val);
                    popNum ++;
                }
                sb.append(val + " ");
                if(popNum == k){
                    break;
                }
                boolean flag = false;
                if(mp.containsKey(dataNode.i-1)){
                    if(mp.get(dataNode.i-1).containsKey(dataNode.j)){
                        flag = true;
                    }
                }
                if(dataNode.i - 1 >= 0 && ! flag) {
                    DataNode d1 = new DataNode(dataNode.i - 1, dataNode.j);
                    priorityQueue.offer(d1);
//                    posOccur[(dataNode.i - 1)][dataNode.j] = true;
                    Map<Integer, Boolean> tmp;
                    if(mp.containsKey(dataNode.i-1)){
                        tmp = mp.get(dataNode.i-1);
                        tmp.put(dataNode.j, true);
                    }else {
                        tmp = new HashMap<>(2);
                        tmp.put(dataNode.j, true);
                    }
                    mp.put(dataNode.i -1, tmp);
                }
                flag = false;
                if(mp.containsKey(dataNode.i)){
                    if(mp.get(dataNode.i).containsKey(dataNode.j-1)){
                        flag = true;
                    }
                }
                if(dataNode.j - 1 >= 0 &&  ! flag) {
                    DataNode d2 = new DataNode(dataNode.i, dataNode.j-1);
                    priorityQueue.offer(d2);
//                    posOccur[dataNode.i][dataNode.j - 1] = true;
                    Map<Integer, Boolean> tmp;
                    if(mp.containsKey(dataNode.i)){
                        tmp = mp.get(dataNode.i);
                        tmp.put(dataNode.j-1, true);
                    }else {
                        tmp = new HashMap<>(2);
                        tmp.put(dataNode.j-1, true);
                    }
                    mp.put(dataNode.i, tmp);
                }
            }else {
                break;
            }
        }
        System.out.println(sb.toString());
    }

}
/*
5 4
1 2 3 4 5
3 5 7 9 11

16 15 14 14

5 30
1 2 3 4 5
3 5 7 9 11

5 -1
1 2 3 4 5
3 5 7 9 11
 */