A21 堆序列 • EH's Blog

题目情境#

题目链接：https://www.luogu.com.cn/problem/P1631 ↗

题目描述

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

输入格式

第一行一个正整数 $N$ ；

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

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

输出格式

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

输入输出样例

输入 #1

3
2 6 6
1 4 8

plaintext

输出 #1

3 6 7

plaintext

说明/提示

对于 $50\%$ 的数据， $N \le 10^3$ 。

对于 $100\%$ 的数据， $1 \le N \le 10^5$ ， $1 \le a_i,b_i \le 10^9$ 。

算法解析#

让 A 中的每个数与 B 中一个数相加，生成 N 个有序队列，

1: A[1] + B[1] ≤ A[2] + B[1] ≤ … ≤ A[N] + B[1] 2: A[1] + B[2] ≤ A[2] + B[2] ≤ … ≤ A[N] + B[2] k: …… N: A[1] + B[N] ≤ A[2] + B[N] ≤ … ≤ A[N] + B[N]

要将这 N 个有序队列进行合并排序：

将这 N 个队列中的第一列元素放入一个小根堆中；
每次取出堆中的最小值。若这个最小值来自于第 k 个队列，就将第 k 个队列的下一个元素放入堆中。

时间复杂度：O(NlogN)。

Python代码实现#

import heapq
n = int(input())
q = []
a = list(map(int, input().split()))
a.insert(0, 0)
b = list(map(int, input().split()))
b.insert(0, 0)
id = [0] * (n + 1)
for i in range(1, n + 1):
    id[i] = 1
    heapq.heappush(q, [a[1] + b[i], i])
for j in range(n):
    print(q[0][0], end=" ")
    i = q[0][1]
    heapq.heappop(q)
    id[i] += 1
    heapq.heappush(q, [a[id[i]] + b[i], i])

python

C++代码实现#

#include <cstdio>
#include <queue>
using namespace std;

const int N=100005;
int a[N],b[N],id[N];
priority_queue<pair<int,int>,
       vector<pair<int,int>>,
       greater<pair<int,int>>>q;
//id[i]: 记录b[i]的搭档的下标
//q: 小根堆，存储<两数和,组的下标>

int main(){
  int n; scanf("%d",&n);
  for(int i=1; i<=n; i++) scanf("%d",&a[i]);
  for(int i=1; i<=n; i++){
    scanf("%d",&b[i]);
    id[i]=1;
    q.push({a[1]+b[i],i});
  }
  while(n--){
    printf("%d ",q.top().first);
    int i=q.top().second; q.pop();
    q.push({a[++id[i]]+b[i],i});
  }
  return 0;
}