Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;

bool checked[100001];
bool selected[100001];
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(false);

  long long int n, k, a, b;
  long long int res = 0;
  int x, y;
  vector<pair<long long int, long long int>> v;
  cin >> n >> k;

  for (int i = 0; i < n; i++) {
    cin >> a >> b;
    v.push_back(make_pair(b, a));
  }

  sort(v.rbegin(), v.rend());

  // 1.
  memset(checked, false, sizeof(checked));
  long long int sum = 0;
  long long int cnt = 0;
  long long int count = 0;
  for (int i = 0; i < n; i++) {
    if (checked[v[i].second] == false) {
      checked[v[i].second] = true;
      cnt++;
    }
    sum += v[i].first;
    count++;
    if (count == k) {
      break;
    }
  }
  res = max(res, (cnt * cnt) + sum);
  x = cnt;
  // cout << cnt << '\n';
  // 2.
  memset(selected, false, sizeof(selected));
  memset(checked, false, sizeof(checked));
  sum = 0;
  cnt = 0;
  count = 0;
  for (int i = 0; i < n; i++) {
    if (checked[v[i].second] == false) {
      checked[v[i].second] = true;
      selected[i] = true;
      sum += v[i].first;
      cnt++;
      count++;
    }
    if (count == k) {
      break;
    }
  }

  for (int i = 0; i < n; i++) {
    if (count == k) {
      break;
    }
    if (selected[i] == false) {
      sum += v[i].first;
      count++;
    }
  }

  res = max(res, sum + (cnt * cnt));
  y = cnt;
  // cout << cnt << '\n';

  for (int i = x + 1; i < y; i++) {
    memset(selected, false, sizeof(selected));
    memset(checked, false, sizeof(checked));
    sum = 0;
    cnt = 0;
    count = 0;

    for (int j = 0; j < n; j++) {
      if (checked[v[j].second] == false) {
        checked[v[j].second] = true;
        selected[j] = true;
        sum += v[j].first;
        cnt++;
        count++;
      }
      if (cnt == i || count == k) {
        break;
      }
    }

    for (int j = 0; j < n; j++) {
      if (count == k) {
        break;
      }
      if (checked[v[j].second] == true && selected[j] == false) {
        sum += v[j].first;
        count++;
      }
    }

    res = max(res, cnt * cnt + sum);
  }
  cout << res << '\n';

  return 0;
}

Fixed submission:
#include <bits/stdc++.h>

using namespace std;

bool checked[100001];
bool selected[100001];
int main(void) {
  cin.tie(0);
  ios::sync_with_stdio(false);

  long long int n, k, a, b;
  long long int res = 0;
  int x, y;
  vector<pair<long long int, long long int>> v;
  cin >> n >> k;

  for (int i = 0; i < n; i++) {
    cin >> a >> b;
    v.push_back(make_pair(b, a));
  }

  sort(v.rbegin(), v.rend());

  // 1.
  memset(checked, false, sizeof(checked));
  long long int sum = 0;
  long long int cnt = 0;
  long long int count = 0;
  for (int i = 0; i < n; i++) {
    if (checked[v[i].second] == false) {
      checked[v[i].second] = true;
      cnt++;
    }
    sum += v[i].first;
    count++;
    if (count == k) {
      break;
    }
  }
  res = max(res, (cnt * cnt) + sum);
  x = cnt;
  // cout << cnt << '\n';
  // 2.
  memset(selected, false, sizeof(selected));
  memset(checked, false, sizeof(checked));
  sum = 0;
  cnt = 0;
  count = 0;
  for (int i = 0; i < n; i++) {
    if (checked[v[i].second] == false) {
      checked[v[i].second] = true;
      selected[i] = true;
      sum += v[i].first;
      cnt++;
      count++;
    }
    if (count == k) {
      break;
    }
  }

  for (int i = 0; i < n; i++) {
    if (count == k) {
      break;
    }
    if (selected[i] == false) {
      sum += v[i].first;
      count++;
    }
  }

  res = max(res, sum + (cnt * cnt));
  y = cnt;
  // cout << cnt << '\n';

  for (int i = x + 1; i < y; i++) {
    memset(selected, false, sizeof(selected));
    memset(checked, false, sizeof(checked));
    sum = 0;
    cnt = 0;
    count = 0;

    for (int j = 0; j < n; j++) {
      if (checked[v[j].second] == false) {
        checked[v[j].second] = true;
        selected[j] = true;
        sum += v[j].first;
        cnt++;
        count++;
      }
      if (cnt == i || count == k) {
        break;
      }
    }

    for (int j = 0; j < n; j++) {
      if (count == k) {
        break;
      }
      if (checked[v[j].second] == true && selected[j] == false) {
        sum += v[j].first;
        count++;
      }
    }

    if (res > cnt * cnt + sum) {
      break;
    }
    res = max(res, cnt * cnt + sum);
  }
  cout << res << '\n';

  return 0;
}