Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <queue>
#include <set>
#include <string>
#include <vector>

using namespace std;

#define int long long
#define rep(i, n) for (int i = 0; i < (n); i++)
#define INF (long long)(1e18)
#define EPS (1e-20)
#define MOD (int)(1e9 + 7)

#define yn(f) ((f) ? "Yes" : "No")
#define YN(f) ((f) ? "YES" : "NO")

#define MAX 210000

int N, K;
int used[MAX];
int temp[MAX], temp2[MAX], temp3[MAX];
priority_queue<int> td[MAX];
priority_queue<int> fuga;
vector<pair<int, int>> hoge;
vector<int> type;
priority_queue<int> hogehoge;

signed main() {
  cout << fixed << setprecision(7);

  set<int> tt;
  set<int> ttt;
  int count = 0;
  int t, d;
  int ans = 0;

  cin >> N >> K;

  for (int i = 0; i < N; i++) {
    cin >> t >> d;
    hoge.push_back(make_pair(d, t));
    tt.insert(t);
  }

  sort(hoge.begin(), hoge.end());

  for (int i = 0; i < K; i++) {
    temp[i] += hoge[N - i - 1].first;
    if (used[hoge[N - i - 1].second] == 0) {
      temp2[i]++;
      used[hoge[N - i - 1].second] = i + 1;
      temp3[i] = hoge[N - i - 1].second;
      count++;
    }
    temp[i + 1] += temp[i];
    temp2[i + 1] += temp2[i];
    // cout<<i<<" temp = "<<temp[i]<<" "<<hoge[N-i-1].first<<"
    // "<<hoge[N-i-1].second<<endl;
  }

  for (int i = N - K - 1; i >= 0; i--) {
    td[hoge[i].second].push(hoge[i].first);
    // cout<<i<<" td "<<hoge[i].second<<" "<<hoge[i].first<<endl;
  }
  // cout<<"XXX "<<endl;

  for (int y : tt) { // cout<<y<<" "<<" "<<tt.size()<<" "<<used[y]<<endl;
    if (used[y])
      continue;
    if (!td[y].empty())
      hogehoge.push(td[y].top());
  }

  // cout<<"YYY "<<endl;
  ans = temp[K - 1] + temp2[K - 1] * temp2[K - 1];

  // for(int i = 0; i < K; i++){
  // cout<<i<<" temp "<<temp[i]<<endl;
  // }
  for (int i = K - 1, xx = 0, tete = tt.size(); i >= 1;
       i--) { // cout<<i<<" "<<temp[i]<<" "<<xx<<endl;
    // cout<<temp2[i-1]<<" "<<temp2[i]+K-i<<" tete "<<tete<<endl;

    if (temp3[i]) {
      xx += temp[i] - temp[i - 1];
    } else {
      xx += hogehoge.top();
      hogehoge.pop();
    }
    ans =
        max(temp[i - 1] + xx + (temp2[i - 1] + K - i) * (temp2[i - 1] + K - i),
            ans);
    // cout<<" end "<<xx<<" temp[i]-> "<<temp[i-1]<<" ans = "<<ans<<endl;
    if (temp2[i - 1] + K - i == tete) {
      break;
    }
  }

  cout << ans << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <queue>
#include <set>
#include <string>
#include <vector>

using namespace std;

#define int long long
#define rep(i, n) for (int i = 0; i < (n); i++)
#define INF (long long)(1e18)
#define EPS (1e-20)
#define MOD (int)(1e9 + 7)

#define yn(f) ((f) ? "Yes" : "No")
#define YN(f) ((f) ? "YES" : "NO")

#define MAX 210000

int N, K;
int used[MAX];
int temp[MAX], temp2[MAX], temp3[MAX];
priority_queue<int> td[MAX];
priority_queue<int> fuga;
vector<pair<int, int>> hoge;
vector<int> type;
priority_queue<int> hogehoge;

signed main() {
  cout << fixed << setprecision(7);

  set<int> tt;
  set<int> ttt;
  int count = 0;
  int t, d;
  int ans = 0;

  cin >> N >> K;

  for (int i = 0; i < N; i++) {
    cin >> t >> d;
    hoge.push_back(make_pair(d, t));
    tt.insert(t);
  }

  sort(hoge.begin(), hoge.end());

  for (int i = 0; i < K; i++) {
    temp[i] += hoge[N - i - 1].first;
    if (used[hoge[N - i - 1].second] == 0) {
      temp2[i]++;
      used[hoge[N - i - 1].second] = i + 1;
      temp3[i] = hoge[N - i - 1].second;
      count++;
    }
    temp[i + 1] += temp[i];
    temp2[i + 1] += temp2[i];
    // cout<<i<<" temp = "<<temp[i]<<" "<<hoge[N-i-1].first<<"
    // "<<hoge[N-i-1].second<<endl;
  }

  for (int i = N - K - 1; i >= 0; i--) {
    td[hoge[i].second].push(hoge[i].first);
    // cout<<i<<" td "<<hoge[i].second<<" "<<hoge[i].first<<endl;
  }
  // cout<<"XXX "<<endl;

  for (int y : tt) { // cout<<y<<" "<<" "<<tt.size()<<" "<<used[y]<<endl;
    if (used[y])
      continue;
    if (!td[y].empty())
      hogehoge.push(td[y].top());
  }

  // cout<<"YYY "<<endl;
  ans = temp[K - 1] + temp2[K - 1] * temp2[K - 1];

  // for(int i = 0; i < K; i++){
  // cout<<i<<" temp "<<temp[i]<<endl;
  // }
  for (int i = K - 1, xx = 0, tete = tt.size(); i >= 1;
       i--) { // cout<<i<<" "<<temp[i]<<" "<<xx<<endl;
    // cout<<temp2[i-1]<<" "<<temp2[i]+K-i<<" tete "<<tete<<endl;

    if (temp3[i]) {
      xx += temp[i] - temp[i - 1];
    } else {
      if (hogehoge.size() == 0)
        break;
      xx += hogehoge.top();
      hogehoge.pop();
    }
    ans =
        max(temp[i - 1] + xx + (temp2[i - 1] + K - i) * (temp2[i - 1] + K - i),
            ans);
    // cout<<" end "<<xx<<" temp[i]-> "<<temp[i-1]<<" ans = "<<ans<<endl;
    if (temp2[i - 1] + K - i == tete) {
      break;
    }
  }

  cout << ans << endl;

  return 0;
}