ABC 116 D - Various Sushi

様々な寿司。

問題概要

N個の寿司に対して、種類と美味しさがそれぞれ与えられる。
美味しさと種類数の2乗を合計したものが満足度となる。
N個中K個の寿司を食べるときの満足度の最大値を求めよ。

atcoder.jp

解法に至るまで

とりあえず美味しさか種類でソートしたい。
種類ボーナスで得られる満足度は種類数の2乗なので、種類ボーナスの方が美味しさより大きければ、すなわち

x^2-(x-1)^2>d[i]

となるなら種類を増やせばよさそう。
種類を増やすかどうかを考えていけば良いから、とりあえず美味しさでソートしてみよう。

解法

とりあえず美味しさでソートして美味しいものを食べていく。

種類を増やせばもっと満足できるからもしれないから、種類を増やしていきたい。
この時、一番おいしくないものを他の種類の寿司と入れ替えたい。
しかし既に1つしか食べていないものを戻すと元も子もないので、2つ以上食べたネタを戻す。
つまり、2つ以上食べたネタで一番美味しくないネタを、食べたことのないネタの中で一番おいしいものと順に入れ替えていけばよい。
美味しさでソートしているので、下図のように、食べたことのあるネタはもちろんスルーして良い。

次に美味しくないネタを見たとき、食べたことのないネタなら入れ替えてみる。

(世界一図を作るのが下手)（図ミスってますけど察してください）

ループの最後でその都度満足度を計算して、高ければ上書きしていけば良い。

ソースコード

実際には種類と美味しさをペアにして美味しさでソート、食べた寿司の種類は配列にメモしていけば良いです。
寿司のネタを入れ替える際、メモの更新、入れ替えた寿司の分いったんスコアをマイナスすることを忘れずに。

    #include<bits/stdc++.h>
    using namespace std;
    typedef long long ll;
     
    int main(){
    	ll n,k;
    	cin>>n>>k;
    	ll t[n],d[n],v[100005];
    	pair<ll,ll> td[n];
    	for(ll i=0;i<n;i++){
    		cin>>t[i]>>d[i];
    		td[i].first = d[i];
    		td[i].second = t[i];
    	}
    	for(ll i=0;i<100005;i++)v[i]=0;
    	sort(td,td+n,greater<pair<ll,ll>>());
    	ll tmp=0,res=0,kind_cnt=0;
    	for(ll i=0;i<k;i++){
    		res += td[i].first;
    		v[td[i].second] ++;
    		if(v[td[i].second]==1)kind_cnt++;
    	}
    	res += kind_cnt*kind_cnt;
    	//cout<<"RES"<<res<<endl;
    	ll cnt=0,sum=res;
    	for(ll i=k-1;i>0;i--){
    		if(v[td[i].second]>1){
    			cnt=-1;
    			for(ll j=k;j<n;j++){
    				if(v[td[j].second]==0){
    					cnt = j;
    					break;
    				}
    			}
    			if(cnt==-1LL)break;
    			sum -= td[i].first; //スコアから、食べるのをやめた寿司の分を引く
    			//cout<<"tdi:"<<td[i].first<<"smu:"<<sum<<endl;
    			v[td[i].second]--; //それぞれのネタを食べた数のメモから、食べるのをやめた分減らす
    			sum+=td[cnt].first; //新たに食べようとしている寿司の分スコアにプラス
    			//cout<<"tdcnt:"<<td[cnt].first<<"smu:"<<sum<<endl;
    			//cout<<"i:"<<i<<"smu:"<<sum<<endl;
    			//cout<<cnt<<endl;
    			kind_cnt++; //種類数を更新
    			v[td[cnt].second]++; //メモの種類ごとに食べたかずを更新
    			sum += (kind_cnt*kind_cnt - (kind_cnt-1)*(kind_cnt-1)); //最後に種類ボーナスを計算
    		}
    		res = max(res,sum); //スコアはそのままで、初めの結果よりスコアが高ければ上書き
    	}
    	cout<<res<<endl;
    }

解説では優先度付きキューを使うことになっていますが、使わなくても間に合います。
（優先度付きキューの方がより高速です）