Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bits/stdc++.h>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <list>
#include <queue>
#include <set>
#include <string>
#include <vector>

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep_r(i, n) for (int i = n - 1; i >= 0; i--)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)
#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define all(x) (x).begin(), (x).end()
#define SZ(x) ((ll)(x).size())
#define bit(n) (1LL << (n))
#define UNIQUE(v) v.erase(unique(v.begin(), v.end()), v.end());

#define INF bit(60)
#define pb push_back
#define mod 1000000007

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

using namespace std;
using uif = uint_fast64_t;
using ll = long long int;

int qp(int a, ll b) {
  int ans = 1;
  do {
    if (b & 1)
      ans = 1ll * ans * a % mod;
    a = 1ll * a * a % mod;
  } while (b >>= 1);
  return ans;
}
int qp(int a, ll b, int mo) {
  int ans = 1;
  do {
    if (b & 1)
      ans = 1ll * ans * a % mo;
    a = 1ll * a * a % mo;
  } while (b >>= 1);
  return ans;
}

int gcd(int a, int b) { return b ? gcd(b, a % b) : a; }
int dx[4] = {1, 0, -1, 0};
int dy[4] = {0, 1, 0, -1};

#define FACSIZE 2002

ll invfac[FACSIZE];
ll fac[FACSIZE];

ll mod_pow(ll, ll);
ll mod_add(ll, ll);
ll mod_mul(ll, ll);
ll mod_div(ll, ll);
void make_fact(ll);
void make_invfact(ll);

int main(void) {
  ll ans = 0;

  using pa = pair<ll, ll>;

  ll n, k;
  cin >> n >> k;
  vector<ll> t(n);
  vector<ll> d(n);

  priority_queue<pa> pq;
  priority_queue<pa, vector<pa>, greater<pa>> pq_remove_cand;

  rep(i, n) {
    cin >> t[i] >> d[i];
    pq.push(make_pair(d[i], t[i]));
  }

  unordered_map<ll, ll> um;
  unordered_map<ll, ll> um_kind;
  ll score = 0;
  ll kind = 0;
  ll score_base = 0;
  rep(i, k) {
    auto q = pq.top();
    // cout << q.first << "," << q.second << endl;
    pq.pop();
    pq_remove_cand.push(q);
    score_base += q.first;
    um[q.second]++;
  }
  kind = um.size();
  // cout << "kind:" << kind << endl;
  score = score_base + kind * kind;
  ans = max(ans, score);
  // cout << "ans:" << ans << endl;
  while (!pq.empty()) {

    auto q = pq.top();
    pq.pop();
    if (um.find(q.second) != end(um))
      continue;
    else {
      // cout << "In:" << q.first << "," << q.second << endl;
      pq_remove_cand.push(q);
      um[q.second]++;
      while (1) {
        auto cq = pq_remove_cand.top();
        pq_remove_cand.pop();
        if (um[cq.second] > 1) {
          // cout << "Out:" << cq.first << "," << cq.second << endl;
          um[cq.second]--;
          score = score + (kind + 1) * (kind + 1) - kind * kind + q.first -
                  cq.first;
          ans = max(ans, score);
          // cout << "ans:" << ans << endl;
          kind++;
          // cout << "kind:" << kind << endl;
          break;
        }
      }
    }
  }

  cout << ans << endl;

  return 0;
}

ll mod_pow(ll x, ll r) {
  if (r == 0)
    return 1;
  else if (r == 1)
    return x % mod;
  else if (r % 2 == 0) {
    ll t = mod_pow(x, r / 2) % mod;
    return t * t % mod;
  } else {
    ll t = mod_pow(x, r / 2) % mod;
    ll k = t * t % mod;
    return k * x % mod;
  }
}
ll mod_add(ll a, ll b) { return (a + b) % mod; }

ll mod_mul(ll a, ll b) { return (a * b) % mod; }

ll mod_div(ll a, ll b) { return mod_mul(a, mod_pow(b, mod - 2)); }

void make_fact(ll n) {
  fac[0] = 1;
  rep(i, n) { fac[i + 1] = mod_mul(fac[i], i + 1); }
}

void make_invfact(ll n) {
  invfac[n] = mod_pow(fac[n], mod - 2);
  for (int i = n - 1; i >= 0; i--) {
    invfac[i] = mod_mul(invfac[i + 1], i + 1);
  }
}

Fixed submission:
#include <algorithm>
#include <bits/stdc++.h>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <list>
#include <queue>
#include <set>
#include <string>
#include <vector>

#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define rep_r(i, n) for (int i = n - 1; i >= 0; i--)
#define rep1(i, n) for (int i = 1; i <= (int)(n); i++)
#define REP(i, n) for (int i = 0, i##_len = (n); i < i##_len; ++i)
#define all(x) (x).begin(), (x).end()
#define SZ(x) ((ll)(x).size())
#define bit(n) (1LL << (n))
#define UNIQUE(v) v.erase(unique(v.begin(), v.end()), v.end());

#define INF bit(60)
#define pb push_back
#define mod 1000000007

template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}

using namespace std;
using uif = uint_fast64_t;
using ll = long long int;

int qp(int a, ll b) {
  int ans = 1;
  do {
    if (b & 1)
      ans = 1ll * ans * a % mod;
    a = 1ll * a * a % mod;
  } while (b >>= 1);
  return ans;
}
int qp(int a, ll b, int mo) {
  int ans = 1;
  do {
    if (b & 1)
      ans = 1ll * ans * a % mo;
    a = 1ll * a * a % mo;
  } while (b >>= 1);
  return ans;
}

int gcd(int a, int b) { return b ? gcd(b, a % b) : a; }
int dx[4] = {1, 0, -1, 0};
int dy[4] = {0, 1, 0, -1};

#define FACSIZE 2002

ll invfac[FACSIZE];
ll fac[FACSIZE];

ll mod_pow(ll, ll);
ll mod_add(ll, ll);
ll mod_mul(ll, ll);
ll mod_div(ll, ll);
void make_fact(ll);
void make_invfact(ll);

int main(void) {
  ll ans = 0;

  using pa = pair<ll, ll>;

  ll n, k;
  cin >> n >> k;
  vector<ll> t(n);
  vector<ll> d(n);

  priority_queue<pa> pq;
  priority_queue<pa, vector<pa>, greater<pa>> pq_remove_cand;

  rep(i, n) {
    cin >> t[i] >> d[i];
    pq.push(make_pair(d[i], t[i]));
  }

  unordered_map<ll, ll> um;
  unordered_map<ll, ll> um_kind;
  ll score = 0;
  ll kind = 0;
  ll score_base = 0;
  rep(i, k) {
    auto q = pq.top();
    // cout << q.first << "," << q.second << endl;
    pq.pop();
    pq_remove_cand.push(q);
    score_base += q.first;
    um[q.second]++;
  }
  kind = um.size();
  // cout << "kind:" << kind << endl;
  score = score_base + kind * kind;
  ans = max(ans, score);
  // cout << "ans:" << ans << endl;
  while (!pq.empty()) {

    auto q = pq.top();
    pq.pop();
    if (um.find(q.second) != end(um))
      continue;
    else {
      // cout << "In:" << q.first << "," << q.second << endl;
      pq_remove_cand.push(q);
      um[q.second]++;
      while (!pq_remove_cand.empty()) {
        auto cq = pq_remove_cand.top();
        pq_remove_cand.pop();
        if (um[cq.second] > 1) {
          // cout << "Out:" << cq.first << "," << cq.second << endl;
          um[cq.second]--;
          score = score + (kind + 1) * (kind + 1) - kind * kind + q.first -
                  cq.first;
          ans = max(ans, score);
          // cout << "ans:" << ans << endl;
          kind++;
          // cout << "kind:" << kind << endl;
          break;
        }
      }
    }
  }

  cout << ans << endl;

  return 0;
}

ll mod_pow(ll x, ll r) {
  if (r == 0)
    return 1;
  else if (r == 1)
    return x % mod;
  else if (r % 2 == 0) {
    ll t = mod_pow(x, r / 2) % mod;
    return t * t % mod;
  } else {
    ll t = mod_pow(x, r / 2) % mod;
    ll k = t * t % mod;
    return k * x % mod;
  }
}
ll mod_add(ll a, ll b) { return (a + b) % mod; }

ll mod_mul(ll a, ll b) { return (a * b) % mod; }

ll mod_div(ll a, ll b) { return mod_mul(a, mod_pow(b, mod - 2)); }

void make_fact(ll n) {
  fac[0] = 1;
  rep(i, n) { fac[i + 1] = mod_mul(fac[i], i + 1); }
}

void make_invfact(ll n) {
  invfac[n] = mod_pow(fac[n], mod - 2);
  for (int i = n - 1; i >= 0; i--) {
    invfac[i] = mod_mul(invfac[i + 1], i + 1);
  }
}