Submission #4398117

---

Source Code Expand

#include <bits/stdc++.h>
using namespace std;

using VS = vector<string>;    using LL = long long;
using VI = vector<int>;       using VVI = vector<VI>;
using PII = pair<int, int>;   using PLL = pair<LL, LL>;
using VL = vector<LL>;        using VVL = vector<VL>;

#define ALL(a)  begin((a)),end((a))
#define RALL(a) (a).rbegin(), (a).rend()
#define SZ(a) int((a).size())
#define SORT(c) sort(ALL((c)))
#define RSORT(c) sort(RALL((c)))
#define UNIQ(c) (c).erase(unique(ALL((c))), end((c)))
#define FOR(i, s, e) for (int(i) = (s); (i) < (e); (i)++)
#define FORR(i, s, e) for (int(i) = (s); (i) > (e); (i)--)
//#pragma GCC optimize ("-O3") 
#ifdef YANG33
#include "mydebug.hpp"
#else
#define DD(x) 
#endif
const int INF = 1e9;                          const LL LINF = 1e16;
const LL MOD = 1000000007;                    const double PI = acos(-1.0);
int DX[8] = { 0, 0, 1, -1, 1, 1, -1, -1 };    int DY[8] = { 1, -1, 0, 0, 1, -1, 1, -1 };

/* -----  2019/02/26  Problem: ABC 116 D / Link: http://abc116.contest.atcoder.jp/tasks/abc116_d  ----- */
/* ------問題------



-----問題ここまで----- */
/* -----解説等-----



----解説ここまで---- */

LL ans = 0LL;

int main() {
	cin.tie(0);
	ios_base::sync_with_stdio(false);

	LL N, K; cin >> N >> K;
	VL t(N), d(N);
	FOR(i, 0, N) {
		cin >> t[i] >> d[i];
		t[i]--;
	}

	VI index(N, 0);
	iota(ALL(index), 0);
	sort(ALL(index),
		[&](int a, int b) {
		return d[a] > d[b];
	});

	VI used(N, 0);
	LL sum = 0;
	LL typecnt = 0;
	priority_queue<LL,VL,greater<LL>>pq;
	FOR(i, 0, K) {
		int idx = index[i];
		sum += d[idx];
		if (used[t[idx]]) {
			pq.push(d[idx]);
		}
		else {
			used[t[idx]] = 1;
			typecnt++;
		}
	}
	ans = sum + typecnt * typecnt;
	DD(de(ans, sum, typecnt));
	FOR(i, K, N) {
		int idx = index[i];
		if (used[t[idx]]) {
			continue;
		}
		if (SZ(pq)) {
			LL swapval = pq.top(); pq.pop();
			DD(de(swapval));
			used[t[idx]] = 1;
			typecnt++;

			sum = sum - swapval + d[idx];
			DD(de(sum))
			ans = max(ans, sum+typecnt*typecnt);
		}
	}

	cout << ans << "\n";

	return 0;
}

Submission Info

Judge Result