给定一个长度为 N 的整数数列:A1, A2, ... , AN。你要重复以下操作 K 次:
每次选择数列中最小的整数(如果最小值不止一个,选择最靠前的),将其删除。
并把与它相邻的整数加上被删除的数值。
输出 K 次操作后的序列。
输入格式
第一行包含两个整数 N 和 K。
第二行包含 N 个整数,A1, A2, ... , AN。
对于 20% 的数据,1 ≤ K < N ≤ 10000。
对于 100% 的数据,1 ≤ K < N ≤ 5 × 1e5,0 ≤ Ai ≤ 1e8。
输出格式
输出 N − K 个整数,中间用一个空格隔开,代表 K 次操作后的序列。
输入样例
5 3
1 4 2 8 7
输出样例
17 7
数据范围与提示
数列变化如下,中括号里的数是当次操作中被选择的数:
[1] 4 2 8 7
5 [2] 8 7
[7] 10 7
17 7
解析:
用优先队列维护每个值和其对应的下标,并且双链表储存整个区间。因为每次取出的最小值,其相邻的元素要加上这个最小值,但是队列中无法取出,所以用 cnt 数组记录每个元素需要增加的值。
每次取出队头的元素,如果这个 cnt 数组中没有记录这个元素,则从链表中删除,并且 cnt 记录其相邻元素所要增加的元素。如果 cnt 中记录了这个元素,则加上 cnt 中记录的增加值,然后放回队列中。
当队列中元素为 n-k 时跳出循环,然后遍历 cnt 将其剩余没有计算的增加值加到队列中剩余的数中。由于剩余元素仍然按照下标的大小排序,所以按照下标输出即可。
我们可知,时间复杂度为线性的。
代码:
#include<bits/stdc++.h>
using namespace std;
const int N=5e5+10;
int n,k;
priority_queue<pair<long long,int>,vector<pair<long long,int>>,greater<pair<long long,int>>>q;
int pre[N],ne[N];
long long cnt[N],a[N],t;
int main(){
scanf("%d%d",&n,&k);
for(int i=1;i<=n;i++){
scanf("%lld",&t);
q.push({t,i});
pre[i]=i-1;
ne[i]=i+1;
}
while(q.size()>n-k){
long long x=q.top().first;
int id=q.top().second;
q.pop();
if(cnt[id]){
q.push({x+cnt[id],id});
cnt[id]=0;
}
else{
int left=pre[id],right=ne[id];
cnt[left]+=x;
cnt[right]+=x;
ne[left]=right;
pre[right]=left;
}
}
while(q.size()){
long long x=q.top().first;
int id=q.top().second;
q.pop();
a[id]=x+cnt[id];
}
for(int i=1;i<=n;i++){
if(a[i]) cout<<a[i]<<" ";
}
return 0;
}