洛谷：P1631 序列合并（优先队列）

原文链接：P1631 序列合并

题目描述

有两个长度都是N的序列A和B，在A和B中各取一个数相加可以得到N^2个和，求这N^2个和中最小的N个。

输入输出格式

输入格式： 第一行一个正整数N；

第二行N个整数Ai,满足Ai<=Ai+1且Ai<=10^9;

第三行N个整数Bi, 满足Bi<=Bi+1且Bi<=10^9.

【数据规模】

对于50%的数据中，满足1<=N<=1000；

对于100%的数据中，满足1<=N<=100000。

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

输入输出样例

输入样例#1： 
3 
2 6 6 
1 4 8 
输出样例#1： 
3 6 7

【解题思路】

1. 首先，对两个序列A和B进行排序，以确保能从最小的数开始组合，从而逐步找到最小的N个和。
2. 使用一个最小堆（优先队列）来存储可能的最小和以及对应的索引。初始时，可以将A中的每个元素与B中的第一个元素相加，并将这些和以及对应的索引（即(A中的索引，B中的索引=0））放入最小堆中。
3. 然后，开始从最小堆中取出元素（这些是当前已知的最小和）。每次取出一个元素后，尝试将对应的A中的元素与B中的下一个元素相加（即如果当前取出的和是由A[i]+B[j]形成的，那么我们将尝试A[i]+B[j+1]），并将新的和以及更新后的索引放回最小堆中。
4. 重复上述步骤，直到找到了N个最小的和为止。注意，为了避免重复计算，当从A[i]+B[j]转向A[i]+B[j+1]时，需要有机制避免再次考虑到A[i]+B[j]这样的组合。

【过程详解】

实际样例来逐步理解是如何工作的。假设我们的输入样例如下：

N = 3

A = [1, 2, 3]

B = [2, 3, 4]

目标是从A和B中任取一个数相加，找出这N^2 = 9N2=9个和中最小的N=3个和。

首先，对A和B进行排序，不过在这个例子中它们已经是有序的。

接下来，初始化优先队列，并将A中的每个元素与B中第一个元素的和放入优先队列中。优先队列中的元素以及它们对应的索引如下：

优先队列: [(3, (0, 0)), (4, (1, 0)), (5, (2, 0))]

解释: 每个元组的形式为(和, (A的索引, B的索引))。

现在，开始从优先队列中依次取出最小元素：

1. 取出(3, (0, 0))，输出和为3，并将(1, 2)（即A[0] + B[1]）放入优先队列。

   • 更新后的优先队列: [(4, (1, 0)), (4, (0, 1)), (5, (2, 0))]

2. 取出(4, (1, 0))，输出和为4，并将(2, 3)（即A[1] + B[1]）放入优先队列。

   • 更新后的优先队列: [(4, (0, 1)), (5, (1, 1)), (5, (2, 0))]

3. 接着，取出(4, (0, 1))，输出和为4。此时，我们应该将(1, 3)（即A[0] + B[2]）放入优先队列。

   • 更新后的优先队列: [(5, (1, 1)), (5, (2, 0)), (5, (0, 2))]

已经找到了最小的3个和，分别是3, 4, 4。这就是最终的输出。

整个过程中，通过贪心的方式确保每次都能取到当前已知的最小和，而通过维护一个优先队列，能有效地管理这些和以及它们对应的索引。

【代码实现】

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<int> A(N), B(N);
    for (int i = 0; i < N; ++i) cin >> A[i];
    for (int i = 0; i < N; ++i) cin >> B[i];

    // 对两个数组进行排序
    sort(A.begin(), A.end());
    sort(B.begin(), B.end());

    // 使用优先队列（最小堆）存储和以及对应的索引
    priority_queue<pair<int, pair<int, int>>, vector<pair<int, pair<int, int>>>, greater<>> pq;

    // 初始将A中每个元素与B[0]的和以及对应索引放入最小堆中
    for (int i = 0; i < N; ++i) {
        pq.push({A[i] + B[0], {i, 0}});
    }

    // 循环N次，每次取出一个当前已知的最小和
    for (int i = 0; i < N; ++i) {
        auto top = pq.top();
        pq.pop();
        cout << top.first << (i < N-1 ? " " : "\n"); // 输出当前的最小和
        int idxA = top.second.first;
        int idxB = top.second.second;

        if (idxB + 1 < N) { // 如果B中还有未使用的元素，将它与当前A中的元素相加，并放入最小堆
            pq.push({A[idxA] + B[idxB + 1], {idxA, idxB + 1}});
        }
    }

    return 0;
}