ABC116 D - Various Sushi

D - Various Sushi

概要

項の数列とが与えられる。この中から項を選び（とで添字共通）、「の項の和(の項に現れる数値の種類数の2乗) 」(と呼ぶ)を最大化せよ。

制約

方針

①先にを降順ソートし、先頭から項を使用してを計算する。
②その後、それ以降の数列を見ていき、の種類数が増えるようなものがあれば、
先に足してあったもののうち最も小さい項で種類数を減らさない項を消去し、上記項を追加した場合のを計算する。

①時点のAと、②の最大値のAを保持しつつ、Aは毎回更新してゆく。
最終的には①のAと②の最大値のAの大きい方を出力する。

感想

ほぼ解説どおり。
消去する項の管理はpriority queueを使わず、種類数を減らさないような項を大きい順にvectorに入れて、pop_backしていった。

解答

Submission #4058221 - AtCoder Beginner Contest 116

#include <bits/stdc++.h>
 
#define FOR(i,a,b) for(int i=(a);i<(b);i++)
#define rep(i,n) FOR(i,0,n)
#define RFOR(i,a,b) for(int i=(a)-1;i>=(b);i--)
#define rrep(i,n) RFOR(i,n,0)
#define fi first
#define se second
 
using namespace std;
 
typedef long long ll;
typedef unsigned long long ull;
typedef pair<ll, ll> pll;
 
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	ll N, K;
	cin >> N >> K;
 
	ll t[N], d[N];
	map<ll, ll> mp1;
	rep(i, N) {
		cin >> t[i] >> d[i];
	}
 
	vector<pll> v;
 
	rep(i, N) {
		v.push_back((pll){d[i], t[i]});
	};
 
	sort(v.begin(), v.end());
	reverse(v.begin(), v.end());
 
	map<ll, ll> mp;
	ll ans = 0;
	ll typ = 0;
	vector<ll> w;
	for(ll i = 0; i < K; i++){
		pll p = v[i];
		auto it = mp.find(p.se);
		if(it == mp.end()){
			mp[p.se] = 0;
			typ++;
		}
 
		mp[p.se]++;
		if(mp[p.se] >= 2) w.push_back(p.fi);
 
		ans += p.fi;
	}
 
	ans += typ * typ;
	// cout << "first:" << ans << endl;
	ll ans2 = ans;
	ll anst = ans;
 
	for(int i = K; i < N; i++){
		pll p = v[i];
		// cout << i << " " << p.se << endl;
		auto it = mp.find(p.se);
		if(it != mp.end()) continue;
 
		mp[p.se] = 1;
 
		if(w.size() == 0) continue;
		ans2 -= w[w.size()-1];
		ans2 += p.fi;
		ans2 -= typ * typ;
		ans2 += (typ+1) * (typ+1);
 
		w.pop_back();
		typ++;
 
		anst = max(anst, ans2);
		// cout << i << ":" << ans << " " << typ << endl;
	}
 
	cout << max(ans, anst) << endl;
}