Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#define ft first
#define sc second
#define lb lower_bound
#define ub upper_bound
#define pb push_back
#define pt(sth) cout << sth << "\n"
#define chmax(a, b)                                                            \
  {                                                                            \
    if (a < b)                                                                 \
      a = b;                                                                   \
  }
#define chmin(a, b)                                                            \
  {                                                                            \
    if (a > b)                                                                 \
      a = b;                                                                   \
  }
#define moC(a, s, b) (a) = ((a)s(b) + MOD) % MOD
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
static const ll INF = 1e18;
static const ll MAX = 1e5 + 7;
static const ll MOD = 1e9 + 7;
ll max(ll a, ll b) { return a > b ? a : b; }
ll min(ll a, ll b) { return a < b ? a : b; }

int main(void) {
  ll N, K;
  cin >> N >> K;
  ll i;
  vector<P> su(N);
  ll ty = 0;
  ll cnt[MAX] = {};

  for (i = 0; i < N; i++) {
    cin >> su[i].sc >> su[i].ft;
    su[i].sc--;
  }
  sort(su.rbegin(), su.rend());

  ll ans = 0;
  ll t = 0;
  priority_queue<P> q;
  for (i = 0; i < K; i++) {
    q.push({-su[i].ft, su[i].sc});
    t += su[i].ft;
    cnt[su[i].sc]++;
    if (cnt[su[i].sc] == 1)
      ty++;
  }
  ans = t;
  ans += ty * ty;

  for (i = K; i < N; i++) {

    while (1) {
      P f = q.top();
      q.pop();
      if (cnt[f.sc] > 1) {
        t -= -f.ft;
        cnt[f.sc]--;
        q.push({-su[i].ft, su[i].sc});
        t += su[i].ft;
        cnt[su[i].sc]++;
        if (cnt[su[i].sc] == 1)
          ty++;
        chmax(ans, t + ty * ty);
        break;
      }
    }
  }

  pt(ans);
}


Fixed submission:
#include <bits/stdc++.h>
#define ft first
#define sc second
#define lb lower_bound
#define ub upper_bound
#define pb push_back
#define pt(sth) cout << sth << "\n"
#define chmax(a, b)                                                            \
  {                                                                            \
    if (a < b)                                                                 \
      a = b;                                                                   \
  }
#define chmin(a, b)                                                            \
  {                                                                            \
    if (a > b)                                                                 \
      a = b;                                                                   \
  }
#define moC(a, s, b) (a) = ((a)s(b) + MOD) % MOD
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
static const ll INF = 1e18;
static const ll MAX = 1e5 + 7;
static const ll MOD = 1e9 + 7;
ll max(ll a, ll b) { return a > b ? a : b; }
ll min(ll a, ll b) { return a < b ? a : b; }

int main(void) {
  ll N, K;
  cin >> N >> K;
  ll i;
  vector<P> su(N);
  ll ty = 0;
  ll cnt[MAX] = {};

  for (i = 0; i < N; i++) {
    cin >> su[i].sc >> su[i].ft;
    su[i].sc--;
  }
  sort(su.rbegin(), su.rend());

  ll ans = 0;
  ll t = 0;
  priority_queue<P> q;
  for (i = 0; i < K; i++) {
    q.push({-su[i].ft, su[i].sc});
    t += su[i].ft;
    cnt[su[i].sc]++;
    if (cnt[su[i].sc] == 1)
      ty++;
  }
  ans = t;
  ans += ty * ty;

  for (i = K; i < N; i++) {
    while (q.size()) {
      P f = q.top();
      q.pop();
      if (cnt[f.sc] > 1) {
        t -= -f.ft;
        cnt[f.sc]--;
        q.push({-su[i].ft, su[i].sc});
        t += su[i].ft;
        cnt[su[i].sc]++;
        if (cnt[su[i].sc] == 1)
          ty++;
        chmax(ans, t + ty * ty);
        break;
      }
    }
  }

  pt(ans);
}