Problem p03148 (C++)

Buggy submission (status: Runtime Error):
//{{{
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <sys/time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using ll = long long;
enum : int { M = (int)1e9 + 7 };
enum : ll { MLL = (ll)1e18L + 9 };
using namespace std;
#ifdef LOCAL
#include "rprint2.hpp"
#else
#define FUNC(name)                                                             \
  template <ostream &out = cout, class... T> void name(T &&...) {}
FUNC(prints)
FUNC(printw) FUNC(printwe) FUNC(printb) FUNC(printbe) FUNC(printd) FUNC(printde)
    FUNC(printdu);
#endif
template <class S, class T> istream &operator>>(istream &in, pair<S, T> &p) {
  return in >> p.first >> p.second;
}
template <class T> istream &operator>>(istream &in, vector<T> &v) {
  for (auto &e : v) {
    in >> e;
  }
  return in;
}
//}}}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> dts(n);
  for (int i = 0; i < n; i++) {
    int t, d;
    cin >> t >> d;
    dts[i] = {d, t};
  }
  sort(dts.rbegin(), dts.rend());
  set<int> eaten;
  ll acc = 0;
  priority_queue<int, vector<int>, greater<int>> pq;
  for (int i = 0; i < k; i++) {
    acc += dts[i].first;
    if (eaten.count(dts[i].second)) {
      pq.push(dts[i].first);
    } else {
      eaten.insert(dts[i].second);
    }
  }
  ll ans = acc + eaten.size() * eaten.size();
  for (int i = k; i < n; i++) {
    printde(eaten);
    int t = dts[i].second;
    if (!eaten.count(t)) {
      eaten.insert(t);
      acc -= pq.top();
      pq.pop();
      acc += dts[i].first;
      ll cand = acc + ll(eaten.size() * eaten.size());
      printde(ans, cand);
      ans = max(ans, cand);
    }
  }
  cout << ans << '\n';
}


Fixed submission:
//{{{
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <sys/time.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using ll = long long;
enum : int { M = (int)1e9 + 7 };
enum : ll { MLL = (ll)1e18L + 9 };
using namespace std;
#ifdef LOCAL
#include "rprint2.hpp"
#else
#define FUNC(name)                                                             \
  template <ostream &out = cout, class... T> void name(T &&...) {}
FUNC(prints)
FUNC(printw) FUNC(printwe) FUNC(printb) FUNC(printbe) FUNC(printd) FUNC(printde)
    FUNC(printdu);
#endif
template <class S, class T> istream &operator>>(istream &in, pair<S, T> &p) {
  return in >> p.first >> p.second;
}
template <class T> istream &operator>>(istream &in, vector<T> &v) {
  for (auto &e : v) {
    in >> e;
  }
  return in;
}
//}}}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> dts(n);
  for (int i = 0; i < n; i++) {
    int t, d;
    cin >> t >> d;
    dts[i] = {d, t};
  }
  sort(dts.rbegin(), dts.rend());
  set<int> eaten;
  ll acc = 0;
  priority_queue<int, vector<int>, greater<int>> pq;
  for (int i = 0; i < k; i++) {
    acc += dts[i].first;
    if (eaten.count(dts[i].second)) {
      pq.push(dts[i].first);
    } else {
      eaten.insert(dts[i].second);
    }
  }
  ll ans = acc + eaten.size() * eaten.size();
  for (int i = k; i < n; i++) {
    if (pq.empty()) {
      break;
    }
    printde(eaten);
    int t = dts[i].second;
    if (!eaten.count(t)) {
      eaten.insert(t);
      acc -= pq.top();
      pq.pop();
      acc += dts[i].first;
      ll cand = acc + ll(eaten.size() * eaten.size());
      printde(ans, cand);
      ans = max(ans, cand);
    }
  }
  cout << ans << '\n';
}