Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cmath>
#include <complex>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <string>
#include <vector>

// output
#define SPBR(w, n) std::cout << (w + 1 == n ? '\n' : ' ');
#define YES cout << "YES" << endl
#define Yes cout << "Yes" << endl
#define NO cout << "NO" << endl
#define No cout << "No" << endl
// utility
#define ALL(i) (i).begin(), (i).end()
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (n)-1; i >= (a); --i)
#define REP(i, n) for (int i = 0; i < int(n); ++i)
#define RREP(i, n) for (int i = int(n) - 1; i >= 0; --i)
#define IN(a, x, b) (a <= x && x < b)
#define OUT(a, x, b) (x < a || b <= x)
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

// type/const
#define int ll
using ll = long long;
using ull = unsigned long long;
using ld = long double;
const int MOD = 1000000007;
/* const int MOD = 998244353; */
const int INF = 1e18;
const double PI = acos(-1);

using namespace std;
struct INIT {
  INIT() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
  }
} INIT;

signed main() {
  int N, K;
  cin >> N >> K;

  vector<int> t(N), d(N);
  REP(i, N) {
    cin >> t[i] >> d[i];
    t[i]--;
  }

  priority_queue<int> q;
  priority_queue<int, vector<int>, greater<int>> eat;
  vector<pair<int, int>> MAX(N, {-INF, -1});
  vector<bool> used(N, false);

  int ret = 0;
  REP(i, N) {
    if (MAX[t[i]].first < d[i]) {
      if (MAX[t[i]].second != -1)
        used[MAX[t[i]].second] = false;
      if (MAX[t[i]].first != -INF)
        ret -= MAX[t[i]].first;
      ret += d[i];
      MAX[t[i]] = {d[i], i};
      used[MAX[t[i]].second] = true;
    }
  }

  int cnt = 0;
  REP(i, N) {
    if (used[i]) {
      cnt++;
      eat.push(d[i]);
    } else
      q.push(d[i]);
  }

  while ((int)eat.size() > K) {
    cnt--;
    ret -= eat.top();
    eat.pop();
  }
  REP(i, K - eat.size()) {
    ret += q.top();
    q.pop();
  }

  int ans = ret + cnt * cnt;
  while (!eat.empty()) {
    cnt--;
    ret -= eat.top();
    eat.pop();
    ret += q.top();
    q.pop();

    chmax(ans, ret + cnt * cnt);
  }

  cout << ans << "\n";

  return 0;
}

Fixed submission:
#include <algorithm>
#include <bitset>
#include <cmath>
#include <complex>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <string>
#include <vector>

// output
#define SPBR(w, n) std::cout << (w + 1 == n ? '\n' : ' ');
#define YES cout << "YES" << endl
#define Yes cout << "Yes" << endl
#define NO cout << "NO" << endl
#define No cout << "No" << endl
// utility
#define ALL(i) (i).begin(), (i).end()
#define FOR(i, a, n) for (int i = (a); i < (n); ++i)
#define RFOR(i, a, n) for (int i = (n)-1; i >= (a); --i)
#define REP(i, n) for (int i = 0; i < int(n); ++i)
#define RREP(i, n) for (int i = int(n) - 1; i >= 0; --i)
#define IN(a, x, b) (a <= x && x < b)
#define OUT(a, x, b) (x < a || b <= x)
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

// type/const
#define int ll
using ll = long long;
using ull = unsigned long long;
using ld = long double;
const int MOD = 1000000007;
/* const int MOD = 998244353; */
const int INF = 1e18;
const double PI = acos(-1);

using namespace std;
struct INIT {
  INIT() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(10);
  }
} INIT;

signed main() {
  int N, K;
  cin >> N >> K;

  vector<int> t(N), d(N);
  REP(i, N) {
    cin >> t[i] >> d[i];
    t[i]--;
  }

  priority_queue<int> q;
  priority_queue<int, vector<int>, greater<int>> eat;
  vector<pair<int, int>> MAX(N, {-INF, -1});
  vector<bool> used(N, false);

  int ret = 0;
  REP(i, N) {
    if (MAX[t[i]].first < d[i]) {
      if (MAX[t[i]].second != -1)
        used[MAX[t[i]].second] = false;
      if (MAX[t[i]].first != -INF)
        ret -= MAX[t[i]].first;
      ret += d[i];
      MAX[t[i]] = {d[i], i};
      used[MAX[t[i]].second] = true;
    }
  }

  int cnt = 0;
  REP(i, N) {
    if (used[i]) {
      cnt++;
      eat.push(d[i]);
    } else
      q.push(d[i]);
  }

  while ((int)eat.size() > K) {
    cnt--;
    ret -= eat.top();
    eat.pop();
  }
  REP(i, K - eat.size()) {
    ret += q.top();
    q.pop();
  }

  int ans = ret + cnt * cnt;
  while (!eat.empty() && !q.empty()) {
    cnt--;
    ret -= eat.top();
    eat.pop();
    ret += q.top();
    q.pop();

    chmax(ans, ret + cnt * cnt);
  }

  cout << ans << "\n";

  return 0;
}