《算法笔记》9.7小节——数据结构专题(2)-＞堆 问题 B: 序列合并

题目描述

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

输入

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

输出

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

样例输入

3
2 6 6
1 4 8

样例输出

3 6 7

提示

建议用最小堆实现。

分析：用最小堆和BFS进行求和。注意用vector模拟堆的时候，下标从0开始。

初始时，两数组下标全部为0。之后堆顶元素出堆，并由堆顶元素进行扩展，两个下标分别+1，如果此时的下标对没有出现过，则将对应的元素之和压入堆中。一直到求出n个和为止。

#include<algorithm>
#include <iostream>
#include  <cstdlib>
#include  <cstring>
#include   <string>
#include   <vector>
#include   <cstdio>
#include    <queue>
#include    <stack>
#include    <ctime>
#include    <cmath>
#include      <map>
#include      <set>
#define INF 0xffffffff
#define db1(x) cout<<#x<<"="<<(x)<<endl
#define db2(x,y) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<endl
#define db3(x,y,z) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<", "<<#z<<"="<<(z)<<endl
#define db4(x,y,z,r) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<", "<<#z<<"="<<(z)<<", "<<#r<<"="<<(r)<<endl
#define db5(x,y,z,r,w) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<", "<<#z<<"="<<(z)<<", "<<#r<<"="<<(r)<<", "<<#w<<"="<<(w)<<endl
using namespace std;

typedef struct node
{
    long long val;
    int i,j;
}node;

void up_heap(vector<node>&heap)
{
    int l=heap.size()-1;
    while(l>0)
    {
        if(heap[l].val<heap[(l-1)/2].val)swap(heap[l],heap[(l-1)/2]),l=(l-1)/2;
        else return;
    }
}

void down_heap(vector<node>& heap)
{
    int l = heap.size();
    int index = 0;

    while (index * 2 + 1 < l) {  // 左子节点的索引为 index * 2 + 1
        int left = index * 2 + 1;
        int right = index * 2 + 2;
        int smallest = index;

        if (left < l && heap[left].val < heap[smallest].val) {
            smallest = left;
        }

        if (right < l && heap[right].val < heap[smallest].val) {
            smallest = right;
        }

        if (smallest != index) {
            swap(heap[index], heap[smallest]);
            index = smallest;
        }
        else {
            return;
        }
    }
}


void heap_pop(vector<node>&heap)
{
    int l=heap.size();
    heap[0]=heap[l-1];
    heap.pop_back();
    down_heap(heap);
}

void update_heap(vector<node>&heap,long long sum,int i,int j)
{
    node temp;temp.val=sum,temp.i=i,temp.j=j;
    heap.push_back(temp);
    up_heap(heap);
}

void get_ans(vector<long long>&ans,vector<node>&heap,int n,int num_a[],int num_b[])
{
    int index_a=0,index_b=0;
    map<long long,int>mp;mp[0]=1;
    node temp;temp.val=num_a[0]+num_b[0],temp.i=0,temp.j=0;
    heap.push_back(temp);
    while(ans.size()<n)
    {
//        db5(ans.size(),heap.size(),heap[0].val,heap[0].i,heap[0].j);

        ans.push_back(heap[0].val);
        int l=heap.size();

//        for(int i=0;i<l;++i)
//        {
//            db4(i,heap[i].val,heap[i].i,heap[i].j);
//        }

        temp=heap[0];heap[0]=heap[l-1];heap.pop_back();
        if(l>0)down_heap(heap);

        index_a=temp.i,index_b=temp.j;
        long long temp_index;
        temp_index=(index_a+1)*1000000+index_b;

//        db4(temp_index,mp[temp_index],index_a,index_b);

        if(!mp[temp_index])mp[temp_index]=1,update_heap(heap,(long long)(num_a[index_a+1]+num_b[index_b]),index_a+1,index_b);

        temp_index=index_a*1000000+(index_b+1);

//        db4(temp_index,mp[temp_index],index_a,index_b);

        if(!mp[temp_index])mp[temp_index]=1,update_heap(heap,(long long)(num_a[index_a]+num_b[index_b+1]),index_a,index_b+1);

//        printf("\n");
    }
}

int main(void)
{
    #ifdef test
    freopen("in.txt","r",stdin);
//    freopen("out.txt","w",stdout);
    clock_t start=clock();
    #endif //test

    int n;scanf("%d",&n);
    int num_a[n+5]={0},num_b[n+5]={0};
    for(int i=0;i<n;++i)scanf("%d",&num_a[i]);
    for(int i=0;i<n;++i)scanf("%d",&num_b[i]);

    vector<node>heap;
    vector<long long>ans;
    get_ans(ans,heap,n,num_a,num_b);

    for(int i=0;i<n;++i)
        printf("%lld ",ans[i]);
    printf("\n");

    #ifdef test
    clockid_t end=clock();
    double endtime=(double)(end-start)/CLOCKS_PER_SEC;
    printf("\n\n\n\n\n");
    cout<<"Total time:"<<endtime<<"s"<<endl;        //s为单位
    cout<<"Total time:"<<endtime*1000<<"ms"<<endl;    //ms为单位
    #endif //test
    return 0;
}