Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include "bits/stdc++.h"
using namespace std;
#define int long long
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define RFOR(i, a, b) for (int i = (b - 1); i >= (a); i--)
#define REP(i, n) for (int i = 0; i < (n); i++)
#define RREP(i, n) for (int i = (n - 1); i >= 0; i--)
#define REP1(i, n) for (int i = 1; i <= (n); i++)
#define RREP1(i, n) for (int i = (n); i >= 1; i--)
#define ALL(a) (a).begin(), (a).end()
#define UNIQUE_SORT(l)                                                         \
  sort(ALL(l));                                                                \
  l.erase(unique(ALL(l)), l.end());
#define CONTAIN(a, b) find(ALL(a), (b)) != (a).end()
#define out(...) printf(__VA_ARGS__)
#if DEBUG
#define debug(...) printf(__VA_ARGS__)
#else
#define debug(...) /* ... */
#endif
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
void _OUT(){};
template <class H, class... T> void _OUT(H &h, T &...t) {
  cout << ' ' << h;
  _OUT(t...);
};
template <class H, class... T> void _P(H &h, T &...t) {
  cout << h;
  _OUT(t...);
};
#define _PN(...)                                                               \
  _P(__VA_ARGS__);                                                             \
  cout << endl;

void solve();
signed main() {
#if DEBUG
  std::ifstream in("input.txt");
  std::cin.rdbuf(in.rdbuf());
#endif
  cin.tie(0);
  ios::sync_with_stdio(false);
  solve();
  return 0;
}

/*================================*/
#if DEBUG
#define SIZE 100
#else
#define SIZE 223450
#endif

int N, K;
vector<int> C(SIZE), Q;

void solve() {
  cin >> N >> K;
  int t, d;
  priority_queue<pair<int, int>> que;
  REP(i, N) {
    cin >> t >> d;
    que.push({d, t});
  }
  int ans = 0, cur = 0, s = 0;
  while (K--) {
    auto [d, t] = que.top();
    que.pop();
    cur += d;
    if (!C[t])
      s++;
    else
      Q.push_back(d);
    C[t] = 1;
  }
  chmax(ans, cur + s * s);
  while (que.size()) {
    auto [d, t] = que.top();
    que.pop();
    if (!C[t]) {
      s++;
      cur -= Q.back();
      cur += d;
      Q.pop_back();
      C[t] = 1;
      chmax(ans, cur + s * s);
    }
  }
  _PN(ans);
}


Fixed submission:
#include "bits/stdc++.h"
using namespace std;
#define int long long
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define RFOR(i, a, b) for (int i = (b - 1); i >= (a); i--)
#define REP(i, n) for (int i = 0; i < (n); i++)
#define RREP(i, n) for (int i = (n - 1); i >= 0; i--)
#define REP1(i, n) for (int i = 1; i <= (n); i++)
#define RREP1(i, n) for (int i = (n); i >= 1; i--)
#define ALL(a) (a).begin(), (a).end()
#define UNIQUE_SORT(l)                                                         \
  sort(ALL(l));                                                                \
  l.erase(unique(ALL(l)), l.end());
#define CONTAIN(a, b) find(ALL(a), (b)) != (a).end()
#define out(...) printf(__VA_ARGS__)
#if DEBUG
#define debug(...) printf(__VA_ARGS__)
#else
#define debug(...) /* ... */
#endif
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
void _OUT(){};
template <class H, class... T> void _OUT(H &h, T &...t) {
  cout << ' ' << h;
  _OUT(t...);
};
template <class H, class... T> void _P(H &h, T &...t) {
  cout << h;
  _OUT(t...);
};
#define _PN(...)                                                               \
  _P(__VA_ARGS__);                                                             \
  cout << endl;

void solve();
signed main() {
#if DEBUG
  std::ifstream in("input.txt");
  std::cin.rdbuf(in.rdbuf());
#endif
  cin.tie(0);
  ios::sync_with_stdio(false);
  solve();
  return 0;
}

/*================================*/
#if DEBUG
#define SIZE 100
#else
#define SIZE 223450
#endif

int N, K;
vector<int> C(SIZE), Q;

void solve() {
  cin >> N >> K;
  int t, d;
  priority_queue<pair<int, int>> que;
  REP(i, N) {
    cin >> t >> d;
    que.push({d, t});
  }
  int ans = 0, cur = 0, s = 0;
  while (K--) {
    auto [d, t] = que.top();
    que.pop();
    cur += d;
    if (!C[t])
      s++;
    else
      Q.push_back(d);
    C[t] = 1;
  }
  chmax(ans, cur + s * s);
  while (que.size() && Q.size()) {
    auto [d, t] = que.top();
    que.pop();
    if (!C[t]) {
      s++;
      cur -= Q.back();
      cur += d;
      Q.pop_back();
      C[t] = 1;
      chmax(ans, cur + s * s);
    }
  }
  _PN(ans);
}