Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <climits>
#include <cmath>
#include <deque>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <utility>

using namespace std;

pair<long long, int> dt[100000];

int main() {
  int N;
  int K;
  cin >> N >> K;
  for (int i = 0; i < N; i++) {
    int t;
    long long d;
    cin >> t >> d;
    dt[i] = make_pair(d, t);
  }

  sort(dt, dt + N);

  long long ans = 0;
  map<int, int> type_counts;
  for (int i = 0; i < K; i++) {
    ans += dt[N - 1 - i].first;
    if (type_counts.find(dt[N - 1 - i].second) != type_counts.end()) {
      type_counts[dt[N - 1 - i].second]++;
    } else {
      type_counts[dt[N - 1 - i].second] = 1;
    }
  }

  long long type_num = type_counts.size();
  ans += type_num * type_num;

  vector<long long> discard_order;
  for (int i = N - K; i < N; i++) {
    if (type_counts[dt[i].second] != 1) {
      discard_order.push_back(dt[i].first);
      type_counts[dt[i].second]--;
    }
  }

  if (discard_order.size() == 0) {
    cout << ans << endl;
  }

  long long tmp_ans = ans;
  int discard_ind = 0;
  for (int i = N - 1 - K; i >= 0; i--) {
    if (type_counts.find(dt[i].second) != type_counts.end()) {
      continue;
    }
    tmp_ans = tmp_ans - type_num * type_num + (type_num + 1) * (type_num + 1) +
              dt[i].first - discard_order[discard_ind];
    ans = max(ans, tmp_ans);
    type_num++;
    discard_ind++;
    type_counts[dt[i].second] = 1;

    if (discard_ind >= discard_order.size()) {
      break;
    }
  }

  cout << ans << endl;

  return 0;
}


Fixed submission:
#include <algorithm>
#include <climits>
#include <cmath>
#include <deque>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <utility>

using namespace std;

pair<long long, int> dt[100000];

int main() {
  int N;
  int K;
  cin >> N >> K;
  for (int i = 0; i < N; i++) {
    int t;
    long long d;
    cin >> t >> d;
    dt[i] = make_pair(d, t);
  }

  sort(dt, dt + N);

  long long ans = 0;
  map<int, int> type_counts;
  for (int i = 0; i < K; i++) {
    ans += dt[N - 1 - i].first;
    if (type_counts.find(dt[N - 1 - i].second) != type_counts.end()) {
      type_counts[dt[N - 1 - i].second]++;
    } else {
      type_counts[dt[N - 1 - i].second] = 1;
    }
  }

  long long type_num = type_counts.size();
  ans += type_num * type_num;

  vector<long long> discard_order;
  for (int i = N - K; i < N; i++) {
    if (type_counts[dt[i].second] != 1) {
      discard_order.push_back(dt[i].first);
      type_counts[dt[i].second]--;
    }
  }

  if (discard_order.size() == 0) {
    cout << ans << endl;
    return 0;
  }

  long long tmp_ans = ans;
  int discard_ind = 0;
  for (int i = N - 1 - K; i >= 0; i--) {
    if (type_counts.find(dt[i].second) != type_counts.end()) {
      continue;
    }
    tmp_ans = tmp_ans - type_num * type_num + (type_num + 1) * (type_num + 1) +
              dt[i].first - discard_order[discard_ind];
    ans = max(ans, tmp_ans);
    type_num++;
    discard_ind++;
    type_counts[dt[i].second] = 1;

    if (discard_ind >= discard_order.size()) {
      break;
    }
  }

  cout << ans << endl;

  return 0;
}