Codeforces Round #629 (Div. 3) F. Make k Equal（前缀和+思维）

题目链接

思路：v.first表示元素的值，second表示元素出现的个数，num【i】代表【1，i-1】的元素的个数前缀和，sum1【i】代表【1，i-1】中元素的个数*值得前缀和，num2【i】和sum2【i】同理求得是【i+1，n】得前缀和，我们知道最终结果肯定是初始元素内得某个数变成了k个，那么我们就枚举假设把当前数变成k个，遍历取min即可，那么对于a【i】来说要把它变成k的操作就是，如果num1【i】+v.second大于等于k的话就计算一下左边的贡献，同理理算右边的贡献，同时不要忘了如果左右单个都不满足的话要取中间的，也就是ans=min(ans,left+right+k-v[i].second)。，left和right的话分别表示把【1，i-1】的元素变成a【i-1】需要的操作和【i+1，n】的元素变成a【i+1】需要的操作。

#include<bits/stdc++.h>
#include<unordered_map>
using namespace std;
typedef long long ll;
const int maxn=2e5+1;
map<ll,ll>p;
vector<pair<ll,ll>>v;
ll k,sum1[maxn],sum2[maxn],num1[maxn],num2[maxn],a[maxn];
int main()
{
	int n;
	scanf("%d %lld",&n,&k);
	for(int i=1;i<=n;++i) scanf("%lld",&a[i]),p[a[i]]++;
	for(auto it=p.begin();it!=p.end();it++)
	v.push_back({it->first,it->second});
	ll t=0,tt=0,ans=1e18;
	for(int i=0;i<v.size();++i)
	{
		num1[i]=t;
		sum1[i]=tt;
		t+=v[i].second;
		tt+=v[i].first*v[i].second;
	}
	t=tt=0;
	for(int i=v.size()-1;i>=0;--i)
	{
		num2[i]=t;
		sum2[i]=tt;
		t+=v[i].second;
		tt+=v[i].first*v[i].second;
	}
	for(int i=0;i<v.size();++i)
	{
		if(v[i].second>=k){
			puts("0");return 0;
		}
		ll left=num1[i]*(v[i].first-1)-sum1[i],right=sum2[i]-num2[i]*(v[i].first+1);
		if(v[i].second+num1[i]>=k) ans=min(ans,left+k-v[i].second);
		if(v[i].second+num2[i]>=k) ans=min(ans,right+k-v[i].second);
		ans=min(ans,left+right+k-v[i].second);
	}
	printf("%lld\n",ans);
}