\(\def \set #1#2{\{ #1 \ \vert \ #2 \}}\)
問題の性質
まずは問題の性質から探る．
- 同じ種類の寿司なら，美味しさが高い物から食べるのが最善．
- 同じ種類を食べるときは，最初の 1個だけは種類ボーナスに影響するが， 2個目以降は影響しない．
簡単な問題
種類ボーナスは無視すると， 全ての寿司をまとめて，おいしい順に並べて選べばよい．
今回の問題
種類ボーナスも考慮すると， 種類毎に類別しておきたい． また，同じ種類であっても，最初の 1個と残りは区別する必要がある．
そこで，美味しさと種類ボーナスが増えるか = \(\ (d, k) \), \(k \in 2\) を寿司ごとに割り当てる． ただし，同じ種類の寿司には，一番おいしい物に \(k = 1\) を割り当てて，それ以外は \(k = 0\) を寿司ごとに割り当てる．
実装
何種類食べたかを全探索する．つまり，\(k = 1\) となる寿司から選ぶ個数を全探索する．
使っている記号，マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"
int main() {
ll n, k;
cin >> n >> k;
vvll ds(n);
rep(i,n){
ll t,d; cin >> t >> d;
--t;
ds[t].emplace_back(d);
}
rep(i,n){
sort(rall(ds[i]));
}
vvll v(2);
rep(j,n) if(ds[j].size()){
v[1].emplace_back(ds[j][0]);
srep(k,1,ds[j].size()) v[0].emplace_back(ds[j][k]);
}
rep(i,2) sort(rall(v[i]));
vll s(2);
stack<ll> q0;
rep(i,min(k,(ll)v[0].size())) q0.push(v[0][i]), s[0] += v[0][i];
s[1] = 0;
ll ans = 0;
rep(i, min(k, (ll)v[1].size())){
s[1] += v[1][i];
while(q0.size() + i+1 > k){
s[0] -= q0.top();
q0.pop();
}
ll x = (i+1)*(i+1);
x += s[1];
x += s[0];
chmax(ans, x);
}
cout << ans << endl;
return 0;
}
