Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define whlie while
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define rep(i, N) for (int i = 0; i < (N); i++)
#define repr(i, N) for (int i = (N)-1; i >= 0; i--)
#define rep1(i, N) for (int i = 1; i <= (N); i++)
#define repr1(i, N) for (int i = (N); i > 0; i--)
#define each(x, v) for (auto &x : v)
#define all(v) (v).begin(), (v).end()
#define sz(v) ((int)(v).size())
#define ini(...)                                                               \
  int __VA_ARGS__;                                                             \
  in(__VA_ARGS__)
#define inl(...)                                                               \
  ll __VA_ARGS__;                                                              \
  in(__VA_ARGS__)
#define ins(...)                                                               \
  string __VA_ARGS__;                                                          \
  in(__VA_ARGS__)
using namespace std;
void solve();
using ll = long long;
using vl = vector<ll>;
using vi = vector<int>;
using vvi = vector<vector<int>>;
constexpr int inf = 1001001001;
constexpr ll infLL = (1LL << 61) - 1;
struct IoSetupNya {
  IoSetupNya() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(15);
    cerr << fixed << setprecision(7);
  }
} iosetupnya;
template <typename T, typename U> inline bool amin(T &x, U y) {
  return (y < x) ? (x = y, true) : false;
}
template <typename T, typename U> inline bool amax(T &x, U y) {
  return (x < y) ? (x = y, true) : false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << p.first << " " << p.second;
  return os;
}
template <typename T, typename U>
istream &operator>>(istream &is, pair<T, U> &p) {
  is >> p.first >> p.second;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  int s = (int)v.size();
  rep(i, s) os << (i ? " " : "") << v[i];
  return os;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &x : v)
    is >> x;
  return is;
}
void in() {}
template <typename T, class... U> void in(T &t, U &...u) {
  cin >> t;
  in(u...);
}
void out() { cout << "\n"; }
template <typename T, class... U> void out(const T &t, const U &...u) {
  cout << t;
  if (sizeof...(u))
    cout << " ";
  out(u...);
}
template <typename T> void die(T x) {
  out(x);
  exit(0);
}
#ifdef NyaanDebug
#include "NyaanDebug.h"
#define trc(...)                                                               \
  do {                                                                         \
    cerr << #__VA_ARGS__ << " = ";                                             \
    dbg_out(__VA_ARGS__);                                                      \
  } while (0)
#define trca(v, N)                                                             \
  do {                                                                         \
    cerr << #v << " = ";                                                       \
    array_out(v, N);                                                           \
    cout << endl;                                                              \
  } while (0)
#else
#define trc(...)
#define trca(...)
int main() { solve(); }
#endif
using P = pair<int, int>;
using vp = vector<P>;
constexpr int MOD = /**/ 1000000007; //*/ 998244353;
/////////

void solve() {
  ini(N, K);
  vp p(N);
  in(p);

  int M, m;
  // 種類数のcnt
  {
    set<int> s;
    each(x, p) s.insert(x.first);
    M = s.size();
  }

  auto f = [](P a, P b) { return a.se > b.se; };
  vl ans(K + 1, -infLL);
  vvi cnt(N + 1);
  set<int> s;

  sort(all(p), f);
  ll cur = 0;
  rep(i, K) {
    cur += p[i].se;
    cnt[p[i].fi].pb(p[i].se);
    s.insert(p[i].fi);
  }

  ans[m = s.size()] = cur;

  priority_queue<P, vector<P>, decltype(f)> Q(f);
  each(i, s) {
    if (cnt[i].size() > 1) {
      sort(all(cnt[i]), greater<int>());
      rep1(j, (int)cnt[i].size() - 1) Q.push(P(i, cnt[i][j]));
    }
  }

  auto g = [](P a, P b) { return a.se < b.se; };
  priority_queue<P, vector<P>, decltype(g)> Q2(g);
  for (int i = K; i < N; i++)
    Q2.push(p[i]);

  for (int i = s.size() + 1; i <= min(M, K); i++) {
    auto x = Q.top();
    Q.pop();
    cur -= x.second;
    P nxt;
    do {
      nxt = Q2.top();
      Q2.pop();
    } while (s.find(nxt.first) != s.end());
    cur += nxt.second;
    s.insert(nxt.first);
    ans[i] = cur;
  }

  ll bns = -inf;
  for (ll i = m; i <= M; i++) {
    amax(bns, ans[i] + i * i);
  }
  out(bns);
}

Fixed submission:
#include <bits/stdc++.h>
#define whlie while
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define rep(i, N) for (int i = 0; i < (N); i++)
#define repr(i, N) for (int i = (N)-1; i >= 0; i--)
#define rep1(i, N) for (int i = 1; i <= (N); i++)
#define repr1(i, N) for (int i = (N); i > 0; i--)
#define each(x, v) for (auto &x : v)
#define all(v) (v).begin(), (v).end()
#define sz(v) ((int)(v).size())
#define ini(...)                                                               \
  int __VA_ARGS__;                                                             \
  in(__VA_ARGS__)
#define inl(...)                                                               \
  ll __VA_ARGS__;                                                              \
  in(__VA_ARGS__)
#define ins(...)                                                               \
  string __VA_ARGS__;                                                          \
  in(__VA_ARGS__)
using namespace std;
void solve();
using ll = long long;
using vl = vector<ll>;
using vi = vector<int>;
using vvi = vector<vector<int>>;
constexpr int inf = 1001001001;
constexpr ll infLL = (1LL << 61) - 1;
struct IoSetupNya {
  IoSetupNya() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(15);
    cerr << fixed << setprecision(7);
  }
} iosetupnya;
template <typename T, typename U> inline bool amin(T &x, U y) {
  return (y < x) ? (x = y, true) : false;
}
template <typename T, typename U> inline bool amax(T &x, U y) {
  return (x < y) ? (x = y, true) : false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << p.first << " " << p.second;
  return os;
}
template <typename T, typename U>
istream &operator>>(istream &is, pair<T, U> &p) {
  is >> p.first >> p.second;
  return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  int s = (int)v.size();
  rep(i, s) os << (i ? " " : "") << v[i];
  return os;
}
template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &x : v)
    is >> x;
  return is;
}
void in() {}
template <typename T, class... U> void in(T &t, U &...u) {
  cin >> t;
  in(u...);
}
void out() { cout << "\n"; }
template <typename T, class... U> void out(const T &t, const U &...u) {
  cout << t;
  if (sizeof...(u))
    cout << " ";
  out(u...);
}
template <typename T> void die(T x) {
  out(x);
  exit(0);
}
#ifdef NyaanDebug
#include "NyaanDebug.h"
#define trc(...)                                                               \
  do {                                                                         \
    cerr << #__VA_ARGS__ << " = ";                                             \
    dbg_out(__VA_ARGS__);                                                      \
  } while (0)
#define trca(v, N)                                                             \
  do {                                                                         \
    cerr << #v << " = ";                                                       \
    array_out(v, N);                                                           \
    cout << endl;                                                              \
  } while (0)
#else
#define trc(...)
#define trca(...)
int main() { solve(); }
#endif
using P = pair<int, int>;
using vp = vector<P>;
constexpr int MOD = /**/ 1000000007; //*/ 998244353;
/////////

void solve() {
  ini(N, K);
  vp p(N);
  in(p);

  int M, m;
  // 種類数のcnt
  {
    set<int> s;
    each(x, p) s.insert(x.first);
    M = s.size();
  }

  auto f = [](P a, P b) { return a.se > b.se; };
  vl ans(K + 1, -infLL);
  vvi cnt(N + 1);
  set<int> s;

  sort(all(p), f);
  ll cur = 0;
  rep(i, K) {
    cur += p[i].se;
    cnt[p[i].fi].pb(p[i].se);
    s.insert(p[i].fi);
  }

  ans[m = s.size()] = cur;

  priority_queue<P, vector<P>, decltype(f)> Q(f);
  each(i, s) {
    if (cnt[i].size() > 1) {
      sort(all(cnt[i]), greater<int>());
      rep1(j, (int)cnt[i].size() - 1) Q.push(P(i, cnt[i][j]));
    }
  }

  auto g = [](P a, P b) { return a.se < b.se; };
  priority_queue<P, vector<P>, decltype(g)> Q2(g);
  for (int i = K; i < N; i++)
    Q2.push(p[i]);

  for (int i = s.size() + 1; i <= min(M, K); i++) {
    auto x = Q.top();
    Q.pop();
    cur -= x.second;
    P nxt;
    do {
      nxt = Q2.top();
      Q2.pop();
    } while (s.find(nxt.first) != s.end());
    cur += nxt.second;
    s.insert(nxt.first);
    ans[i] = cur;
  }

  ll bns = -inf;
  for (ll i = m; i <= min(M, K); i++) {
    amax(bns, ans[i] + i * i);
  }
  out(bns);
}