洛谷 P1631 序列合并

题意简述

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

题解思路

大根堆，先存入n个和，再比较大小，改变堆中元素。

代码

#include <cstdio>
#include <algorithm>
using namespace std;
int n, a[100001], b[100001], ans[100001];
int cnt, heap[100001];
void up(int x);
void down(int x);
void push(int x);
void pop();
int main()
{
    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]);
    for (int i = 1; i <= n; ++i)
        push(a[1] + b[i]);
    for (int i = 2; i <= n; ++i)
        for (int j = 1; j <= n; ++j)
            if (a[i] + b[j] < heap[1])
            {
                pop();
                push(a[i] + b[j]);
            }
            else break;
    for (int i = 1; i <= n; ++i)
    {
        ans[n - i + 1] = heap[1];
        pop();
    }
    for (int i = 1; i <= n - 1; ++i)
        printf("%d ", ans[i]);
    printf("%d\n", ans[n]);
    return 0;
}
void up(int x)
{
    if (x == 1) return;
    if (heap[x / 2] < heap[x])
    {
        swap(heap[x / 2], heap[x]);
        up(x / 2);
    }
}
void down(int x)
{
    x *= 2;
    if (x > cnt) return;
    if (x < cnt && heap[x + 1] > heap[x]) x++;
    if (heap[x / 2] < heap[x])
    {
        swap(heap[x / 2], heap[x]);
        down(x);
    }
}
void push(int x)
{
    heap[++cnt] = x;
    up(cnt);
}
void pop()
{
    heap[1] = heap[cnt --];
    down(1);
}