Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cstdio>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <vector>

using namespace std;

typedef long long ll;

int t[100005];
ll d[100005];
vector<pair<ll, int>> vec;
int mp[100005];

struct sushi {
  int ti;
  ll di;
  bool operator>(const sushi &n) const { return di > n.di; }
} eat;

int main() {

  priority_queue<sushi, vector<sushi>, greater<sushi>> pque;

  int N, K;

  scanf("%d %d", &N, &K);

  int i;

  for (i = 0; i < N; i++) {
    scanf("%d %lld", &t[i], &d[i]);
    vec.push_back(make_pair(d[i], t[i]));
  }

  sort(vec.begin(), vec.end());
  reverse(vec.begin(), vec.end());

  ll sum = 0ll;
  ll x = 0ll;
  int pos = K;

  for (i = 0; i < K; i++) {
    sum += vec[i].first;
    mp[vec[i].second]++;
    eat.di = vec[i].first;
    eat.ti = vec[i].second;
    pque.push(eat);
  }

  for (i = 1; i <= N; i++) {
    if (mp[i] > 0) {
      x += 1ll;
    }
  }

  ll ans = sum + (x * x);

  sushi tmp, tmp2;

  while (1) {

    if (pos >= N) {
      break;
    }

    tmp = pque.top();
    pque.pop();

    if (mp[tmp.ti] == 1) {
      continue;
    }

    for (i = pos; i < N; i++) {
      if (mp[vec[i].second] == 0) {
        break;
      }
    }

    if (i < N) {

      ll score = (sum - tmp.di) + (vec[i].first) + ((x + 1ll) * (x + 1ll));

      tmp2.di = vec[i].first;
      tmp2.ti = vec[i].second;

      pque.push(tmp2);
      x += 1ll;
      sum = sum - tmp.di + vec[i].first;
      mp[tmp.ti]--;
      mp[vec[i].second]++;
      ans = max(ans, score);
      pos = i;
    } else {
      break;
    }
  }

  printf("%lld\n", ans);

  return 0;
}

Fixed submission:
#include <algorithm>
#include <cstdio>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <vector>

using namespace std;

typedef long long ll;

int t[100005];
ll d[100005];
vector<pair<ll, int>> vec;
int mp[100005];

struct sushi {
  int ti;
  ll di;
  bool operator>(const sushi &n) const { return di > n.di; }
} eat;

int main() {

  priority_queue<sushi, vector<sushi>, greater<sushi>> pque;

  int N, K;

  scanf("%d %d", &N, &K);

  int i;

  for (i = 0; i < N; i++) {
    scanf("%d %lld", &t[i], &d[i]);
    vec.push_back(make_pair(d[i], t[i]));
  }

  sort(vec.begin(), vec.end());
  reverse(vec.begin(), vec.end());

  ll sum = 0ll;
  ll x = 0ll;
  int pos = K;

  for (i = 0; i < K; i++) {
    sum += vec[i].first;
    mp[vec[i].second]++;
    eat.di = vec[i].first;
    eat.ti = vec[i].second;
    pque.push(eat);
  }

  for (i = 1; i <= N; i++) {
    if (mp[i] > 0) {
      x += 1ll;
    }
  }

  ll ans = sum + (x * x);

  sushi tmp, tmp2;

  while (1) {

    if (pos >= N || pque.empty()) {
      break;
    }

    tmp = pque.top();
    pque.pop();

    if (mp[tmp.ti] == 1) {
      continue;
    }

    for (i = pos; i < N; i++) {
      if (mp[vec[i].second] == 0) {
        break;
      }
    }

    if (i < N) {

      ll score = (sum - tmp.di) + (vec[i].first) + ((x + 1ll) * (x + 1ll));

      tmp2.di = vec[i].first;
      tmp2.ti = vec[i].second;

      pque.push(tmp2);
      x += 1ll;
      sum = sum - tmp.di + vec[i].first;
      mp[tmp.ti]--;
      mp[vec[i].second]++;
      ans = max(ans, score);
      pos = i;
    } else {
      break;
    }
  }

  printf("%lld\n", ans);

  return 0;
}