OI队分糖果加强版

OI队分糖果加强版

Description
L老师买了一包糖，其中有m颗糖果分给OI队N个小朋友 每一个小朋友都希望自己得到一定数量的糖果数，如果没有达到他的希望值ai ​，小朋友就会生气，可以认为他生气的程度等于他少得到的糖果数量的平方值。 比如，DZC想要得到32个糖果，但是只得到了29个。他少了3个，所以他的生气值是9。、 现在糖果数不足以满足所有小朋友的希望。所以我们应该采取最优的分配，使得最后小朋友们的生气值的和最小。

Input
输入数据共n+1行。 第一行两个整数m,n。 接下来n行，每行一个整数，第i+1行的整数表示第i个小朋友希望值ai
1≤n≤10⁵ 1≤m≤2×10⁹ 结果不超过2^64-1

Output
输出数据共一行。 一行一个整数，表示最小的总生气指数。

Examples
input
10 4
4
5
2
3
output
4
提示
共10颗糖果，共4人，给每一个同学他所需要的糖果数减1，也就是依次给3,4,1,2颗，这样的话每人少一颗，每人的生气指数就是1^2=1，44个人就是1×4=4，答案4是最优方案

①暴力版代码，会TLE

#include <bits/stdc++.h>
#define F(a,b,c)for(int a=b;a<=c;a++)
using namespace std;
bool cmp(long long x,long long y){return x>y;}
long long m,n,a[1000010],maxx,qwq=1,ans;
int main(){
	scanf("%lld%lld",&m,&n);
	for(int i=1;i<=n;i++)scanf("%lld",&a[i]);
	sort(a+1,a+n+1,cmp);
	maxx=a[1];
	while(m)
		if(a[qwq]!=maxx)maxx=a[1],qwq=1;
		else a[qwq++]--,m--;
	for(int i=1;i<=n;i++)ans+=a[i]*a[i];
	printf("%lld",ans);
	return 0;
}

②AC代码

#include <bits/stdc++.h>
using namespace std;
typedef long long l;
l n,m,a[100001],b[100001],c=1,ans;
bool cmp(l x,l y){return x>y;}
int main(){
	cin>>m>>n;
	for(l i=1;i<=n;i++)cin>>a[i];
	sort(a+1,a+n+1,cmp);
	for(l i=1;i<=n;i++){
		if(m>=i*(a[i]-a[i+1]))
		m-=i*(a[i]-a[i+1]);
		else{
			for(int j=1;j<=i;j++)a[j]=a[i]-m/i;
			m%=i;
			for(int j=1;j<=m;j++)a[j]--;
			break;
		}
	}
	for(l i=1;i<=n;i++)ans+=a[i]*a[i];
	cout<<ans;
	return 0;
}