标题Codeforces Round #629 (Div. 3)F. Make k Equal (前后缀和+讨论)
下图来自某博客园大佬,我认为写的很赞。
下面是AC代码
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll N=2e5+5,inf=1e15;
ll a[N],n,k,pre[N],suf[N],ans=inf;
int main(){
cin>>n>>k;
for(int i=1;i<=n;i++) cin>>a[i];
sort(a+1,a+n+1);//排序
for(int i=1;i<=n;i++) pre[i]=pre[i-1]+a[i];//前缀
for(int i=n;i>=1;i--) suf[i]=suf[i+1]+a[i];//后缀
for(int i=1,j;i<=n;i++)//枚举左端i
{
j=i;
while(j+1<=n&&a[j+1]==a[i]) j++;//找到右端j
if(j-i+1>=k){ //如果相同个数已经大于等于k个
ans=0;
break;
}
ll tmp=(a[i]-1)*(i-1)-pre[i-1]+suf[j+1]-(n-j)*(a[i]+1)+k-(j-i+1);//两边都取.
ans=min(ans,tmp);//取最小值
if(j>=k){//特判只取左边
ans=min(ans,(a[i]-1)*(i-1)-pre[i-1]+k-(j-i+1));
}
if(n-i+1>=k)//只取右边
ans=min(ans,suf[j+1]-(a[i]+1)*(n-j)+k-(j-i+1));
i=j;//更新
}
cout<<ans<<endl;
return 0;
}