Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define fi first
#define se second
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define rep2(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define rev(i, n) for (int(i) = (n)-1; (i) >= 0; (i)--)
#define rev2(i, a, b) for (int(i) = (a)-1; (i) >= (b); (i)--)
#define rng(a) a.begin(), a.end()
#define rrng(a) a.rbegin(), a.rend()
#define pb push_back
#define eb emplace_back
#define yn                                                                     \
  { puts("Yes"); }                                                             \
  else {                                                                       \
    puts("No");                                                                \
  }
using namespace std;
using ll = long long;
using P = pair<int, int>;
using Pll = pair<long long, long long>;
template <class T> using Vec = vector<T>;
template <class T> using VVec = vector<vector<T>>;
long long TEN(int i) { return !i ? 1 : TEN(i - 1) * 10LL; }
template <class T, class U> bool chmin(T &a, const U &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T, class U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}

struct IoSetup {
  IoSetup() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
  }
} iosetup;

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

int N, K;
int t[100005], d[100005];

int main() {
  cin >> N >> K;
  rep(i, N) cin >> t[i] >> d[i];

  set<int> st;
  int syu = 0;
  rep(i, N) st.insert(t[i]);
  syu = st.size();

  rep(i, N) t[i]--;
  VVec<P> ho(N);
  rep(i, N) ho[t[i]].eb(d[i], 0);
  rep(i, N) sort(rrng(ho[i]));
  rep(i, N) if (ho[i].size()) ho[i][0].se = 1;
  priority_queue<P> Q;
  rep(i, N) rep(j, ho[i].size()) Q.push({ho[i][j].fi, ho[i][j].se});

  /*
  while (!Q.empty()) {
          auto [oi, mask] = Q.top();
          Q.pop();
          cout << oi << "  " << mask << endl;
  }
  */

  ll ans = 0;
  int cnt = 0;
  Vec<int> ko0;
  rep(i, K) {
    auto [oi, mask] = Q.top();
    Q.pop();
    cnt += mask;
    if (!mask)
      ko0.pb(oi);
    ans += oi;
  }
  ans += (ll)cnt * cnt;

  sort(rrng(ko0));
  ll before_ans = ans;
  rep2(i, cnt + 1, syu + 1) {
    ll tans = before_ans;
    tans -= ko0[ko0.size() - 1];
    ko0.pop_back();
    int oi, mask = 0;
    while (!mask) {
      tie(oi, mask) = Q.top();
      Q.pop();
    }
    tans += oi;
    tans -= (ll)(i - 1) * (i - 1);
    tans += (ll)i * i;
    chmax(ans, tans);
    before_ans = tans;
  }

  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define fi first
#define se second
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define rep2(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define rev(i, n) for (int(i) = (n)-1; (i) >= 0; (i)--)
#define rev2(i, a, b) for (int(i) = (a)-1; (i) >= (b); (i)--)
#define rng(a) a.begin(), a.end()
#define rrng(a) a.rbegin(), a.rend()
#define pb push_back
#define eb emplace_back
#define yn                                                                     \
  { puts("Yes"); }                                                             \
  else {                                                                       \
    puts("No");                                                                \
  }
using namespace std;
using ll = long long;
using P = pair<int, int>;
using Pll = pair<long long, long long>;
template <class T> using Vec = vector<T>;
template <class T> using VVec = vector<vector<T>>;
long long TEN(int i) { return !i ? 1 : TEN(i - 1) * 10LL; }
template <class T, class U> bool chmin(T &a, const U &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T, class U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}

struct IoSetup {
  IoSetup() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
  }
} iosetup;

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

int N, K;
int t[100005], d[100005];

int main() {
  cin >> N >> K;
  rep(i, N) cin >> t[i] >> d[i];

  set<int> st;
  int syu = 0;
  rep(i, N) st.insert(t[i]);
  syu = st.size();

  rep(i, N) t[i]--;
  VVec<P> ho(N);
  rep(i, N) ho[t[i]].eb(d[i], 0);
  rep(i, N) sort(rrng(ho[i]));
  rep(i, N) if (ho[i].size()) ho[i][0].se = 1;
  priority_queue<P> Q;
  rep(i, N) rep(j, ho[i].size()) Q.push({ho[i][j].fi, ho[i][j].se});

  /*
  while (!Q.empty()) {
          auto [oi, mask] = Q.top();
          Q.pop();
          cout << oi << "  " << mask << endl;
  }
  */

  ll ans = 0;
  int cnt = 0;
  Vec<int> ko0;
  rep(i, K) {
    auto [oi, mask] = Q.top();
    Q.pop();
    cnt += mask;
    if (!mask)
      ko0.pb(oi);
    ans += oi;
  }
  ans += (ll)cnt * cnt;

  sort(rrng(ko0));
  ll before_ans = ans;
  rep2(i, cnt + 1, syu + 1) {
    ll tans = before_ans;
    if (ko0.size() == 0)
      break;
    tans -= ko0[ko0.size() - 1];
    ko0.pop_back();
    int oi, mask = 0;
    while (!mask) {
      tie(oi, mask) = Q.top();
      Q.pop();
    }
    tans += oi;
    tans -= (ll)(i - 1) * (i - 1);
    tans += (ll)i * i;
    chmax(ans, tans);
    before_ans = tans;
  }

  cout << ans << endl;
  return 0;
}