题解：洛谷 P1631 序列合并

【题目来源】

洛谷：P1631 序列合并 - 洛谷

【题目描述】

有两个长度为 N 的单调不降序列 \(A,B\)，在 \(A,B\) 中各取一个数相加可以得到 \(N^2\) 个和，求这 \(N^2\) 个和中最小的 \(N\) 个。

【输入】

第一行一个正整数 \(N\)；

第二行 \(N\) 个整数 \(A_{1\dots N}\)。

第三行 \(N\) 个整数 \(B_{1\dots N}\)。

【输出】

一行 \(N\) 个整数，从小到大表示这 \(N\) 个最小的和。

【输入样例】

3
2 6 6
1 4 8

【输出样例】

3 6 7

【核心思想】

1. 问题分析：给定两个长度为 \(N\) 的单调不降序列 \(A\) 和 \(B\)，从 \(A\)、\(B\) 中各取一个数相加，共有 \(N^2\) 个和，需要输出其中最小的 \(N\) 个和（按从小到大排列）。直接枚举所有 \(N^2\) 对并排序会超时，需利用序列的单调性和优先队列进行优化。

2. 算法选择：

   • 优先队列（小根堆） + 多路归并：将问题视为 \(N\) 个有序链表的归并。对于 \(B\) 中的每个元素 \(B_i\)，将其与 \(A\) 中的元素依次配对，形成的序列 \(A_1 + B_i, A_2 + B_i, \dots, A_N + B_i\) 是单调递增的。初始时，每个链表只取第一个元素（即 \(A_1 + B_i\)）入堆。每次从堆中弹出最小的和，假设该和来自 \(B_i\) 与 \(A_j\) 的配对，则将该 \(B_i\) 与 \(A_{j+1}\) 的和（即下一个可能的最小值）压入堆中。重复 \(N\) 次即可得到最小的 \(N\) 个和。
   • 复杂度：堆中始终维护 \(N\) 个元素，每次弹出和压入操作 \(O(\log N)\)，总复杂度 \(O(N \log N)\)。

3. 关键步骤：

   • 读入与排序：输入 \(N\) 以及数组 \(A[1..N]\)、\(B[1..N]\)（题目保证已单调不降，无需额外排序）。
   • 初始化：用数组 idx[i] 记录 \(B_i\) 当前已匹配到 \(A\) 的哪个下标，初始均为 1。将 \(A[1] + B_i\) 以及对应的 \(B\) 的下标 \(i\) 组成二元组 \((sum, i)\) 放入小根堆。
   • 循环输出：重复 \(N\) 次：
     • 取出堆顶元素 \((sum, i)\)，输出 \(sum\)。
     • 将 \(idx[i]\) 增加 1，然后将 \(A[idx[i]] + B[i]\) 入堆（若 \(idx[i] \le N\)）。
   • 输出：按顺序输出每次取出的 \(sum\)。

4. 时间/空间复杂度：

   • 时间复杂度：\(O(N \log N)\)，每次堆操作 \(O(\log N)\)，共 \(N\) 次。
   • 空间复杂度：\(O(N)\)，存储数组和堆。

5. 堆 + 多路归并的核心思想：

   • 将两数之和转化为有序链表的归并：固定 \(B\) 的每个元素，它与 \(A\) 的配对结果构成一个长度为 \(N\) 的递增序列。于是 \(N\) 个序列的归并问题即为求所有序列前 \(N\) 小的元素。
   • 利用优先队列维护当前最小值：每次从所有序列的当前指针位置中选出最小的和，输出后将该序列指针后移一位，并压入新的候选值。这保证了每次取出的都是全局当前最小。
   • 有序性的依赖：算法正确性依赖于 \(A\) 和 \(B\) 均单调不降，从而每个链表内部有序，使得每次后移指针产生的值是当前序列的下一个最小值。
   • 适用于多路归并、多序列求前 \(k\) 小元素和等场景。

【算法标签】

普及+ #堆排序

【代码详解】

#include <bits/stdc++.h>
using namespace std;

const int MAX_N = 100005;  // 定义数组最大长度
int n;                     // 数组长度
int a[MAX_N], b[MAX_N];    // 存储两个输入数组
int id[MAX_N];             // id[i]记录b[i]当前配对的a数组下标

// 定义小根堆，存储<两数和, b数组下标>
priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> q;

int main()
{
    // 输入数组长度
    cin >> n;

    // 输入数组a
    for (int i = 1; i <= n; i++)
        cin >> a[i];

    // 输入数组b
    for (int i = 1; i <= n; i++)
        cin >> b[i];

    // 初始化id数组和优先队列
    for (int i = 1; i <= n; i++)
    {
        id[i] = 1;                          // 初始每个b[i]与a[1]配对
        q.push({a[1] + b[i], i});           // 将初始配对和存入队列
    }

    // 输出前n小的两数和
    while (n--)
    {
        // 获取当前最小的两数和
        cout << q.top().first << " ";

        // 获取对应的b数组下标
        int i = q.top().second;
        q.pop();

        // 将该b[i]与a的下一个元素配对
        id[i]++;
        q.push({a[id[i]] + b[i], i});
    }

    return 0;
}

【运行结果】

3
2 6 6
1 4 8
3 6 7