Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; ++i)
#define rrep(i, n) for (int i = n - 1; i >= 0; --i)
#define fi first
#define se second
#define sz(x) (int)(x).size()
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
using namespace std;
using ll = long long;
using P = pair<int, int>;
template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> vector<T> make_vec(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_vec(size_t a, Ts... ts) {
  return vector<decltype(make_vec<T>(ts...))>(a, make_vec<T>(ts...));
}
template <typename T, typename U, typename... V>
typename enable_if<is_same<T, U>::value>::type fill_v(U &u, const V... v) {
  u = U(v...);
}
template <typename T, typename U, typename... V>
typename enable_if<!is_same<T, U>::value>::type fill_v(U &u, const V... v) {
  for (auto &e : u)
    fill_v<T>(e, v...);
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  int n, k;
  cin >> n >> k;
  priority_queue<P> pq1;
  rep(i, n) {
    int t, d;
    cin >> t >> d;
    pq1.emplace(d, t);
  }
  ll now = 0;
  set<int> st;
  priority_queue<int, vector<int>, greater<int>> pq2;
  rep(i, k) {
    P p = pq1.top();
    pq1.pop();
    if (st.count(p.se)) {
      pq2.push(p.fi);
    } else {
      st.insert(p.se);
    }
    now += p.fi;
  }
  ll v = st.size();
  ll ans = now + v * v;
  while (!pq1.empty()) {
    P p = pq1.top();
    pq1.pop();
    if (st.count(p.se))
      continue;
    now -= pq2.top();
    pq2.pop();
    st.insert(p.se);
    now += p.fi;
    v++;
    chmax(ans, now + v * v);
  }
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; ++i)
#define rrep(i, n) for (int i = n - 1; i >= 0; --i)
#define fi first
#define se second
#define sz(x) (int)(x).size()
#define all(v) v.begin(), v.end()
#define rall(v) v.rbegin(), v.rend()
using namespace std;
using ll = long long;
using P = pair<int, int>;
template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (b < a) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> vector<T> make_vec(size_t a) { return vector<T>(a); }
template <typename T, typename... Ts> auto make_vec(size_t a, Ts... ts) {
  return vector<decltype(make_vec<T>(ts...))>(a, make_vec<T>(ts...));
}
template <typename T, typename U, typename... V>
typename enable_if<is_same<T, U>::value>::type fill_v(U &u, const V... v) {
  u = U(v...);
}
template <typename T, typename U, typename... V>
typename enable_if<!is_same<T, U>::value>::type fill_v(U &u, const V... v) {
  for (auto &e : u)
    fill_v<T>(e, v...);
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  int n, k;
  cin >> n >> k;
  priority_queue<P> pq1;
  rep(i, n) {
    int t, d;
    cin >> t >> d;
    pq1.emplace(d, t);
  }
  ll now = 0;
  set<int> st;
  priority_queue<int, vector<int>, greater<int>> pq2;
  rep(i, k) {
    P p = pq1.top();
    pq1.pop();
    if (st.count(p.se)) {
      pq2.push(p.fi);
    } else {
      st.insert(p.se);
    }
    now += p.fi;
  }
  ll v = st.size();
  ll ans = now + v * v;
  while (!pq1.empty() && !pq2.empty()) {
    P p = pq1.top();
    pq1.pop();
    if (st.count(p.se))
      continue;
    now -= pq2.top();
    pq2.pop();
    st.insert(p.se);
    now += p.fi;
    v++;
    chmax(ans, now + v * v);
  }
  cout << ans << endl;
}