兩個有序陣列間相加和的Topk問題
阿新 • • 發佈:2021-10-18
給定兩個有序陣列arr1和arr2,再給定一個整數k,返回來自arr1和arr2的兩個數相加和最大的前k個,兩個數必須分別來自兩個陣列
心之所向,素履以往 生如逆旅,一葦以航import java.util.*; public class Main { private static final int[] DX = {0, -1}; private static final int[] DY = {-1, 0}; private static String genKey(int x, int y) { return x + "_" + y; } private static int[] solve(int[] a, int[] b, int k) { if (a == null || a.length == 0 || b == null || b.length == 0) { return new int[0]; } int n = a.length; int[] ret = new int[k]; Set<String> visited = new HashSet<>(); PriorityQueue<Node> queue = new PriorityQueue<>(new Comparator<Node>() { @Override public int compare(Node o1, Node o2) { return Integer.compare(o2.value, o1.value); } }); queue.offer(new Node(a[n - 1] + b[n - 1], n - 1, n - 1)); visited.add(genKey(n - 1, n - 1)); while (!queue.isEmpty() && k > 0) { Node poll = queue.poll(); ret[--k] = poll.value; for (int i = 0; i < 2; ++i) { int x = poll.x + DX[i]; int y = poll.y + DY[i]; if (x >= 0 && y >= 0 && !visited.contains(genKey(x, y))) { queue.offer(new Node(a[x] + b[y], x, y)); visited.add(genKey(x, y)); } } } return ret; } private static void print(int[] ret) { if (ret == null || ret.length == 0) { System.out.println(); return; } System.out.print(ret[ret.length - 1]); for (int i = ret.length - 2; i >= 0; --i) { System.out.print(" " + ret[i]); } System.out.println(); } public static void main(String[] args) { Scanner in = new Scanner(System.in); while (in.hasNext()) { int n = in.nextInt(); int k = in.nextInt(); int[] a = new int[n]; for (int i = 0; i < n; ++i) { a[i] = in.nextInt(); } int[] b = new int[n]; for (int i = 0; i < n; ++i) { b[i] = in.nextInt(); } int[] ret = solve(a, b, k); print(ret); } } } class Node { int value; int x; int y; public Node(int value, int x, int y) { this.value = value; this.x = x; this.y = y; } }