Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cassert>
#include <cmath>
#include <deque>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unistd.h>
#include <unordered_map>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<vvl> vvvl;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<bool> vb;
typedef vector<vector<bool>> vvb;
typedef vector<string> vs;
typedef vector<vector<string>> vvs;
typedef vector<char> vc;
typedef vector<vector<char>> vvc;
typedef vector<double> vd;
typedef vector<vd> vvd;
typedef vector<vvd> vvvd;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;
typedef vector<P> vp;
typedef vector<PL> vpl;
typedef vector<vector<P>> vvp;
typedef vector<vector<PL>> vvpl;
const int INF = 1001001001;
const ll LINF = 1e16;
const double pi = 3.1415926535897932;
const string endstr = "\n";
#define FOR(i, a, b) for (ll i = (a); i < b; i++)
#define REP(i, n) for (ll i = 0; i < n; i++)
#define FORMAP(it, m) for (auto it = m.begin(); it != m.end(); it++)
#define ff first
#define ss second
#define pb push_back

template <typename T> T gcd(T a, T b) { return (a == 0) ? b : gcd(b % a, a); }
template <typename T> T lcm(T a, T b) { return a / gcd(a, b) * b; }

bool p_comp_fs(const PL p1, const PL p2) { return p1.first < p2.first; };
bool p_comp_fg(const PL p1, const PL p2) { return p1.first > p2.first; };
bool p_comp_ss(const PL p1, const PL p2) { return p1.second < p2.second; };
bool p_comp_sg(const PL p1, const PL p2) { return p1.second > p2.second; };
template <typename T> vector<T> uniquen(vector<T> vec) {
  vec.erase(unique(vec.begin(), vec.end()), vec.end());
  return vec;
}

int main() {
  ll N, K;
  cin >> N >> K;
  vpl dt(N);
  REP(i, N) {
    ll t, d;
    cin >> t >> d;
    dt[i] = PL(d, t);
  }
  sort(dt.begin(), dt.end(), p_comp_fg);
  set<ll> unq_t;
  multimap<ll, ll> dup;
  ll deli = 0;
  REP(i, K) {
    ll d = dt[i].ff, t = dt[i].ss;
    deli += d;
    if (unq_t.count(t) > 0) {
      dup.insert(make_pair(d, t));
    } else {
      unq_t.insert(t);
    }
  }
  ll x = unq_t.size();
  ll ans = deli + x * x;
  REP(i, N) {
    ll d = dt[i].ff, t = dt[i].ss;
    if (unq_t.count(t) > 0)
      continue;
    unq_t.insert(t);
    x++;
    deli += d - dup.begin()->ff;
    ans = max(ans, deli + x * x);
    dup.erase(dup.begin());
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cassert>
#include <cmath>
#include <deque>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unistd.h>
#include <unordered_map>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<vvl> vvvl;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<bool> vb;
typedef vector<vector<bool>> vvb;
typedef vector<string> vs;
typedef vector<vector<string>> vvs;
typedef vector<char> vc;
typedef vector<vector<char>> vvc;
typedef vector<double> vd;
typedef vector<vd> vvd;
typedef vector<vvd> vvvd;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;
typedef vector<P> vp;
typedef vector<PL> vpl;
typedef vector<vector<P>> vvp;
typedef vector<vector<PL>> vvpl;
const int INF = 1001001001;
const ll LINF = 1e16;
const double pi = 3.1415926535897932;
const string endstr = "\n";
#define FOR(i, a, b) for (ll i = (a); i < b; i++)
#define REP(i, n) for (ll i = 0; i < n; i++)
#define FORMAP(it, m) for (auto it = m.begin(); it != m.end(); it++)
#define ff first
#define ss second
#define pb push_back

template <typename T> T gcd(T a, T b) { return (a == 0) ? b : gcd(b % a, a); }
template <typename T> T lcm(T a, T b) { return a / gcd(a, b) * b; }

bool p_comp_fs(const PL p1, const PL p2) { return p1.first < p2.first; };
bool p_comp_fg(const PL p1, const PL p2) { return p1.first > p2.first; };
bool p_comp_ss(const PL p1, const PL p2) { return p1.second < p2.second; };
bool p_comp_sg(const PL p1, const PL p2) { return p1.second > p2.second; };
template <typename T> vector<T> uniquen(vector<T> vec) {
  vec.erase(unique(vec.begin(), vec.end()), vec.end());
  return vec;
}

int main() {
  ll N, K;
  cin >> N >> K;
  vpl dt(N);
  REP(i, N) {
    ll t, d;
    cin >> t >> d;
    dt[i] = PL(d, t);
  }
  sort(dt.begin(), dt.end(), p_comp_fg);
  set<ll> unq_t;
  multimap<ll, ll> dup;
  ll deli = 0;
  REP(i, K) {
    ll d = dt[i].ff, t = dt[i].ss;
    deli += d;
    if (unq_t.count(t) > 0) {
      dup.insert(make_pair(d, t));
    } else {
      unq_t.insert(t);
    }
  }
  ll x = unq_t.size();
  ll ans = deli + x * x;
  REP(i, N) {
    ll d = dt[i].ff, t = dt[i].ss;
    if (unq_t.count(t) > 0)
      continue;
    if (dup.size() < 1)
      break;
    unq_t.insert(t);
    x++;
    deli += d - dup.begin()->ff;
    ans = max(ans, deli + x * x);
    dup.erase(dup.begin());
  }
  cout << ans << endl;
  return 0;
}