Atcoder beginner contest 116D Various Sushi

https://atcoder.jp/contests/abc116/tasks/abc116_d
题意：有n个寿司，每个有一个品种，一个美味值，选k个，总美味值等于选的品种数的平方加美味值的和，求最大的总美味值。
分析：选寿司时应该美味值的优先级高。先选出美味值最高的k个，计算他们的总美味值，再依次增加选取种类的数量，每次更新答案。
思路：这个题可以用贪心。第一种情况是把美味值最大的k个加起来，第二种是取的种类数越多越好。在种类数增多的过程中也要更新答案。注意long long。

#include<cstdio>
#include<string>
#include<cstring>
#include<utility>
#include<cmath>
#include<map>
#include<queue>
#include<set>
#include<algorithm>
#include<vector>
#include<iostream>
#define ll long long
#define pii pair<int,int>
#define mp make_pair
#define fi first
#define se second
#define inf 0x7fffffff
using namespace std;
pii a[100010];
int b[100010],c[100010];
int main()
{
	int i,j,k,n,m,x,y,h=0;
	scanf("%d%d",&n,&m);
	for(i=0;i<n;i++)
	{
		scanf("%d%d",&a[i].se,&a[i].fi);
	}
	sort(a,a+n);
	reverse(a,a+n);
	ll ans=0,an;
	x=0;
	for(i=0;i<m;i++)
	{
		ans+=1ll*a[i].fi;
		if(b[a[i].se]==0)
		{
			x++;
			b[a[i].se]=1;
		}
		else
		{
			c[h]=i;
			h++;
		}
	}
	ans+=1ll*x*x;
	an=ans;
	for(i=m;i<n;i++)
	{
		if(b[a[i].se]==0)
		{
			if(h==0)
			{
				break;
			}
			b[a[i].se]=1;
			h--;
			x++;
			ans-=1ll*a[c[h]].fi;
			ans+=1ll*a[i].fi;
			ans+=2ll*x-1ll;
			an=max(an,ans);
		}
	}
	printf("%lld",an);
}