Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, ll> P;
#define FOR(i, N) for (i = 0; i < N; i++)
int N, K;
vector<P> t;

bool pairCompareFirst(const P &firstElof, const P &secondElof) {
  return firstElof.first > secondElof.first;
}

bool pairCompareSecond(const P &firstElof, const P &secondElof) {
  return firstElof.second > secondElof.second;
}

int main(void) {
  cin >> N >> K;
  t.resize(N);
  // k：食べる種類,k_max：食べる種類の最大値
  int i, k = 0, k_max = 0;
  int num[N + 1];
  // 種類iの最大か否か
  bool kind_max[N];

  vector<ll> base(N, 0);
  ll sum = 0;
  ll max = 0;

  FOR(i, N + 1) num[i] = 0;
  FOR(i, N) { cin >> t[i].first >> t[i].second; }
  sort(t.begin(), t.end(), pairCompareSecond);

  FOR(i, N) {
    k_max += num[t[i].first] == 0;
    kind_max[i] = (num[t[i].first] == 0);
    if (i < K) {
      sum += t[i].second;
      k += num[t[i].first] == 0;
    }
    num[t[i].first]++;
  }
  base[k - 1] = sum;
  max = sum + (ll)k * k;

  int EM_index = K - 1, NM_index = K;
  for (i = k; i <= k_max && i <= K; i++) {
    // 交換可能な場所まで移動
    while (kind_max[EM_index]) {
      EM_index--;
      if (EM_index < 0)
        break;
    }
    while (!kind_max[NM_index]) {
      NM_index++;
      if (NM_index >= N)
        break;
    }
    // cout << EM_index << " " << NM_index << endl;
    base[i] = base[i - 1] - t[EM_index].second + t[NM_index].second;
    if (base[i] + (ll)(i + 1) * (i + 1) > max)
      max = base[i] + (ll)(i + 1) * (i + 1);
    swap(t[EM_index], t[NM_index]);
    swap(kind_max[EM_index], kind_max[NM_index]);
    if (EM_index != 0)
      EM_index--;
    else
      break;
    if (NM_index != N - 1)
      NM_index++;
    else
      break;
  }
  // cout << endl;
  // FOR(i,N) cout << t[i].first << " " << t[i].second << " " << kind_max[i] <<
  // endl;
  cout << max << endl;

  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, ll> P;
#define FOR(i, N) for (i = 0; i < N; i++)
int N, K;
vector<P> t;

bool pairCompareFirst(const P &firstElof, const P &secondElof) {
  return firstElof.first > secondElof.first;
}

bool pairCompareSecond(const P &firstElof, const P &secondElof) {
  return firstElof.second > secondElof.second;
}

int main(void) {
  cin >> N >> K;
  t.resize(N);
  // k：食べる種類,k_max：食べる種類の最大値
  int i, k = 0, k_max = 0;
  int num[N + 1];
  // 種類iの最大か否か
  bool kind_max[N];

  vector<ll> base(N, 0);
  ll sum = 0;
  ll max = 0;

  FOR(i, N + 1) num[i] = 0;
  FOR(i, N) { cin >> t[i].first >> t[i].second; }
  sort(t.begin(), t.end(), pairCompareSecond);

  FOR(i, N) {
    k_max += num[t[i].first] == 0;
    kind_max[i] = (num[t[i].first] == 0);
    if (i < K) {
      sum += t[i].second;
      k += num[t[i].first] == 0;
    }
    num[t[i].first]++;
  }
  base[k - 1] = sum;
  max = sum + (ll)k * k;

  int EM_index = K - 1, NM_index = K;
  for (i = k; i < k_max && i < K; i++) {
    // 交換可能な場所まで移動
    while (kind_max[EM_index]) {
      EM_index--;
      if (EM_index < 0)
        break;
    }
    while (!kind_max[NM_index]) {
      NM_index++;
      if (NM_index >= N)
        break;
    }
    // cout << EM_index << " " << NM_index << endl;
    base[i] = base[i - 1] - t[EM_index].second + t[NM_index].second;
    if (base[i] + (ll)(i + 1) * (i + 1) > max)
      max = base[i] + (ll)(i + 1) * (i + 1);
    swap(t[EM_index], t[NM_index]);
    swap(kind_max[EM_index], kind_max[NM_index]);
    if (EM_index != 0)
      EM_index--;
    else
      break;
    if (NM_index != N - 1)
      NM_index++;
    else
      break;
  }
  // cout << endl;
  // FOR(i,N) cout << t[i].first << " " << t[i].second << " " << kind_max[i] <<
  // endl;
  cout << max << endl;

  return 0;
}