算法笔记练习 9.7 堆 问题 B: 序列合并 - 超级详细的思路讲解

算法笔记练习 题解合集

本题链接

题目

题目描述
有两个长度都为N的序列A和B，在A和B中各取一个数相加可以得到N²个和，求这N²个和中最小的N个。

输入
第一行一个正整数N（1 <= N <= 100000）。
第二行N个整数Ai，满足Ai <= Ai+1且Ai <= 109
第三行N个整数Bi，满足Bi <= Bi+1且Bi <= 109

输出
输出仅有一行，包含N个整数，从小到大输出这N个最小的和，相邻数字之间用空格隔开。

样例输入

3
2 6 6
1 4 8

样例输出

3 6 7

思路

因为只需要求 N² 个和中最小的 N 个，所以可以在计算 N² 个和的过程中维护一个大小为 N 的大顶堆。注意只有大顶堆才能维护最小的 N 个数，因为如果在计算过程中出现一个新的数，它小于现有 N 个数中最大的数，就应该把最大的那个数用新的数替换掉。

---

具体步骤：

1. 用数组a接收序列 A 中的所有数；
2. 初始化一个大小为 N 的大顶堆heap，令input接收序列 B 的第一个数。数组a中的所有数全都加上input即为heap中的初始元素（这是维护的开始，我们已经成生成了 N 个结果，并且显然这 N 个数是目前最小的 N 个数）；
3. 令input继续接收序列 B，每当input接收新的序列 B 中的数字，就让input和序列 A 中的每一个数字相加，生成新的结果sum：
   a. 若sum小于heap[1]（堆顶），则令heap[1] = sum，并调整堆；
   b. 若sum大于或等于heap[1]，直接让input接收下一个数字，因为序列 A 是非递减排列的，如果当前的结果都不小于堆顶，那后面的结果不可能小于堆顶。（这一步能显著节约时间）
4. 进行堆排序，输出答案。

代码

#include <cstdio>
#include <algorithm>
#define MAX 100010

int a[MAX], heap[MAX];

void downAdjust(int low, int high) {
	int i = low, j = i * 2;
	while (j <= high) {
		if (j + 1 <= high && heap[j + 1] > heap[j])
			++j;
		if (heap[j] > heap[i]) {
			std::swap(heap[j], heap[i]);
			i = j;
			j = 2 * i;
		} else
			break; 
	} 
} 
void createHeap(int n) {
	for (int i = n / 2; i >= 1; --i)
		downAdjust(i, n);
}
void heapSort(int n) {
	for (int i = n; i > 1; --i) {
		std::swap(heap[1], heap[i]);
		downAdjust(1, i - 1); 
	} 
} 

int main() {
	int n, input;
	while (scanf("%d", &n) != EOF) {
		for (int i = 1; i <= n; ++i) { 
			scanf("%d", &a[i]);
			heap[i] = a[i];
		}
		scanf("%d", &input);
		for (int i = 1; i <= n; ++i)
			heap[i] += input; 
		createHeap(n);
		for (int i = 2; i <= n; ++i) {
			scanf("%d", &input);
			for (int j = 1; j <= n; ++j) {
				if (input + a[j] < heap[1]) { 
					heap[1] = input + a[j];
					downAdjust(1, n);
				} else
					break;
			} 
		}
		heapSort(n);
		for (int i = 1; i <= n; ++i)
			printf("%s%d", i == 1 ? "" : " ", heap[i]);
		putchar('\n'); 
	} 
	return 0;
}