POJ 2442 - Sequence - [小顶堆][优先队列]
题目链接:http://poj.org/problem?id=2442
Time Limit: 6000MS Memory Limit: 65536K
Description
Given m sequences, each contains n non-negative integer. Now we may select one number from each sequence to form a sequence with m integers. It's clear that we may get n ^ m this kind of sequences. Then we can calculate the sum of numbers in each sequence, and get n ^ m values. What we need is the smallest n sums. Could you help us?
Input
The first line is an integer T, which shows the number of test cases, and then T test cases follow. The first line of each case contains two integers m, n (0 < m <= 100, 0 < n <= 2000). The following m lines indicate the m sequence respectively. No integer in the sequence is greater than 10000.
Output
For each test case, print a line with the smallest n sums in increasing order, which is separated by a space.
Sample Input
1
2 3
1 2 3
2 2 3
Sample Output
3 3 4
题意:
给出 $m$ 个长度为 $n$ 的序列,在每一个序列中挑选一个数求和,可知有 $n^m$ 个结果,要求给出这些结果中前 $n$ 小的。
题解(参考《算法竞赛进阶指南》):
首先考虑当 $M=2$ 时的做法,记这两个序列为 $a,b$,先对所有的序列进行升序排序,并且用两个指针 $p,q$ 分别指向这两个序列的头部 $a[0],b[0]$,显然此时就是第 $1$ 小的和 $S_1 = a[0] + b[0]$。
那么,第 $2$ 小的和 $S_2 = \min (a[1] + b[0], a[0] + b[1])$,
如果第 $2$ 小和是 $S_2 = a[1] + b[0]$,那么此时竞争的第 $3$ 小的候选者,除了上面已经存在的 $a[0] + b[1]$ 外还应当再增加 $a[2] + b[0], a[1] + b[1]$,也就是序列一的指针 $ptr_1++$ 或者序列二的指针 $ptr_2++$。
不难想到,可以用一个小顶堆(或者STL的优先队列)来维护所有候选者,不断地扔进去新晋的候选者,再取出堆顶作为答案之一,再由该堆顶求出新的候选者插入堆中,再取出堆顶,反复如此……
候选方案产生方式如图:
不过,有一点需要注意,例如我们在上述例子中已得 $S_2 = a[1] + b[0]$,那么入堆两个新的候选者之后堆中有三个节点: $a[0] + b[1], a[2] + b[0], a[1] + b[1]$;此时如果 $S_3 = a[0] + b[1]$,会发现又一次产生了候选者 $a[1] + b[1]$,这样就产生了重复,具体表现就是在上图中,就是绝大部分候选方案从 $(0,0)$ 出发可以有多条路径到达。重复方案多次入堆显然是影响正确性的,因此要避免这种情况发生。
不妨增加限定,如果当前选中的方案所产生新的候选者是 $ptr_2++$,那么以后都只能 $ptr_2++$,不能再回到 $ptr_1++$。换句话说,$a[0]+b[0]$ 要走到任何候选方案 $a[i]+a[j]$,必须先移动 $ptr_1 = 0 \sim i$,再移动 $ptr_2 = 0 \sim j$,使得到达备选方案 $a[i]+a[j]$ 的路径的唯一性,这样即可避免产生同一个候选方案重复入队的情况。
增加限定条件后的候选方案产生方式如图(不难看出,已经变为了一棵树):
考虑到添加了该限定条件后,是否影响到算法的正确性:
考虑原算法在选定 $a[i] + b[j]$ 成为第 $k$ 小之后,原本会产生 $a[i+1] + b[j]$ 和 $a[i] + b[j+1]$ 两个新的候选者。那么增加限定条件后,是否会发生本来第 $k+1$ 小应当是 $a[i+1] + b[j]$ 但是现在却没有生成该候选方案的情况?
根据限定条件,可知 $a[i] + b[j]$ 该方案成为第 $k$ 小,必然是由于 $a[i] + b[j-1]$ 产生了它,往前依次类推必然是在某一时刻选择了 $a[i] + b[0]$,而 $a[i] + b[0]$ 会产生两个候选者 $a[i+1] + b[0]$ 和 $a[i] + b[1]$,由此可知选中方案 $a[i] + b[j]$,只要 $j \ge 1$,那么此时队列里必然曾经出现过 $a[i+1] + b[0]$。
因此,如果说选定 $a[i] + b[j]$ 成为第 $k$ 小同时堆中没有 $a[i+1] + b[j]$,那么此时堆中必然存在 $a[i+1] +b[0],a[i+1] +b[1], \cdots ,a[i+1] +b[j-1]$ 中的某一个。显然,存在比 $a[i+1] + b[j]$ 还小的方案,肯定不会选到 $a[i+1] + b[j]$,因此增加该限定条件不影响算法正确性。
时间复杂度:
由于我们每次获取并删除一个堆顶,最多往堆中插入两个新的元素,因此每一次堆的大小最多增加 $1$,又因为最多只有 $N$ 次出堆操作且堆初始为空,因此堆的规模最大为 $O(n)$。
因此每次push和pop都是 $O(\log n)$ 的复杂度,而最多做 $O(n)$ 次push和pop,因此 $O(n \log n)$ 就能求得两个序列的前 $n$ 小的和。
而对于 $M>2$ 的情况,可以先求出第一个序列和第二个序列的前 $n$ 小的和,作为一个新序列再去和第三个序列求前 $n$ 小的和,以此类推总的时间复杂度为 $O(mn \log n)$。
AC代码:
手写二叉堆版本(500ms):
#include<cstdio> #include<iostream> #include<vector> #include<cstring> #include<algorithm> using namespace std; typedef pair<int,int> pii; const int maxm=100+5; const int maxn=2000+5; int m,n; vector<int> a[maxm]; struct Plan{ pii ptr; bool last; //记录本方案是否是由ptr2++得到的 int sum; Plan(){} Plan(int i,int j,pii _ptr,bool _last) { ptr=_ptr; last=_last; sum=a[i][ptr.first]+a[j][ptr.second]; } }init; struct Heap { int sz; Plan heap[3*maxn]; void up(int now) { while(now>1) { int par=now>>1; if(heap[now].sum<heap[par].sum) //子节点小于父节点,不满足小顶堆性质 { swap(heap[par],heap[now]); now=par; } else break; } } void push(const Plan &x) //插入权值为x的节点 { heap[++sz]=x; up(sz); } inline Plan top(){return heap[1];} void down(int now) { while((now<<1)<=sz) { int nxt=now<<1; if(nxt+1<=sz && heap[nxt+1].sum<heap[nxt].sum) nxt++; //取左右子节点中较小的 if(heap[now].sum>heap[nxt].sum) //子节点小于父节点,不满足小顶堆性质 { swap(heap[now],heap[nxt]); now=nxt; } else break; } } void pop() //移除堆顶 { heap[1]=heap[sz--]; down(1); } void del(int p) //删除存储在数组下标为p位置的节点 { heap[p]=heap[sz--]; up(p), down(p); } inline void clr(){sz=0;} }h; int main() { int T; cin>>T; while(T--) { scanf("%d%d",&m,&n); for(int i=1;i<=m;i++) { a[i].clear(); for(int j=1,x;j<=n;j++) { scanf("%d",&x); a[i].push_back(x); } sort(a[i].begin(),a[i].end()); } int i=1; for(int j=2;j<=m;j++,i^=1) { init=Plan(i,j,make_pair(0,0),0); h.clr(); h.push(init); a[i^1].clear(); while(h.sz) { Plan now=h.top(); h.pop(); a[i^1].push_back(now.sum); if(a[i^1].size()>=n) break; if(!now.last && now.ptr.first<n-1) h.push(Plan(i,j,make_pair(now.ptr.first+1,now.ptr.second),0)); if(now.ptr.second<n-1) h.push(Plan(i,j,make_pair(now.ptr.first,now.ptr.second+1),1)); } } for(int k=0;k<a[i].size();k++) printf("%d ",a[i][k]); printf("\n"); } }
优先队列版本(563ms):
#include<cstdio> #include<iostream> #include<vector> #include<queue> #include<cstring> #include<algorithm> using namespace std; typedef pair<int,int> pii; const int maxm=100+5; const int maxn=2000+5; int m,n; vector<int> a[maxm]; struct Plan{ pii ptr; bool last; //记录本方案是否是由ptr2++得到的 int sum; Plan(){} Plan(int i,int j,pii _ptr,bool _last) { ptr=_ptr; last=_last; sum=a[i][ptr.first]+a[j][ptr.second]; } bool operator<(const Plan& oth)const{return sum>oth.sum;} }init; priority_queue<Plan> Q; int main() { int T; cin>>T; while(T--) { scanf("%d%d",&m,&n); for(int i=1;i<=m;i++) { a[i].clear(); for(int j=1,x;j<=n;j++) { scanf("%d",&x); a[i].push_back(x); } sort(a[i].begin(),a[i].end()); } int i=1; for(int j=2;j<=m;j++,i^=1) { init=Plan(i,j,make_pair(0,0),0); while(!Q.empty()) Q.pop(); Q.push(init); a[i^1].clear(); while(!Q.empty()) { Plan now=Q.top(); Q.pop(); a[i^1].push_back(now.sum); if(a[i^1].size()>=n) break; if(!now.last && now.ptr.first<n-1) Q.push(Plan(i,j,make_pair(now.ptr.first+1,now.ptr.second),0)); if(now.ptr.second<n-1) Q.push(Plan(i,j,make_pair(now.ptr.first,now.ptr.second+1),1)); } } for(int k=0;k<a[i].size();k++) printf("%d ",a[i][k]); printf("\n"); } }