Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define rep(i, n) FOR(i, 0, n)
#define rFOR(i, a, b) for (ll i = a - 1; i >= b; i--)
#define rrep(i, a) rFOR(i, a, 0)
#define pb push_back
using namespace std;
using ll = long long;
using ld = long double;
typedef pair<ll, ll> P;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<P> vP;
typedef vector<char> vc;
typedef vector<vc> vvc;
const ll MOD = 1000000007;
const ll MOD2 = 998244353;
const ld PI = acos(-1);
const ll INF = 1e18;
struct edge {
  ll to, cost;
};

template <typename T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int N, K;
  cin >> N >> K;
  vP p(N);
  rep(i, N) { cin >> p[i].second >> p[i].first; }
  sort(p.begin(), p.end(), greater<P>());
  set<int> type;
  multiset<int> kaburi;
  unsigned long long res = 0, ans = 0;
  rep(i, N) {
    if (i < K) {
      if (type.count(p[i].second)) {
        kaburi.insert(p[i].first);
        res += p[i].first;
      } else {
        type.insert(p[i].second);
        res += p[i].first;
      }
    } else {
      if (type.count(p[i].second)) {
        continue;
      } else {
        int x = *kaburi.begin();
        res -= x;
        kaburi.erase(kaburi.find(x));
        res += p[i].first;
        type.insert(p[i].second);
      }
    }
    chmax(ans, res + type.size() * type.size());
  }
  cout << ans << '\n';
} /*
 overflow checked?
 corner case checked?
 boundary checked?
 not TLE in worst case checked?
 */

Fixed submission:
// #define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
#define FOR(i, a, b) for (ll i = a; i < b; i++)
#define rep(i, n) FOR(i, 0, n)
#define rFOR(i, a, b) for (ll i = a - 1; i >= b; i--)
#define rrep(i, a) rFOR(i, a, 0)
#define pb push_back
using namespace std;
using ll = long long;
using ld = long double;
typedef pair<ll, ll> P;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<P> vP;
typedef vector<char> vc;
typedef vector<vc> vvc;
const ll MOD = 1000000007;
const ll MOD2 = 998244353;
const ld PI = acos(-1);
const ll INF = 1e18;
struct edge {
  ll to, cost;
};

template <typename T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int N, K;
  cin >> N >> K;
  vP p(N);
  rep(i, N) { cin >> p[i].second >> p[i].first; }
  sort(p.begin(), p.end(), greater<P>());
  set<int> type;
  multiset<int> kaburi;
  unsigned long long res = 0, ans = 0;
  rep(i, N) {
    if (i < K) {
      if (type.count(p[i].second)) {
        kaburi.insert(p[i].first);
        res += p[i].first;
      } else {
        type.insert(p[i].second);
        res += p[i].first;
      }
    } else {
      if (type.count(p[i].second)) {
        continue;
      } else {
        int x = *kaburi.begin();
        res -= x;
        if (kaburi.size() == 0)
          continue;
        kaburi.erase(kaburi.find(x));
        res += p[i].first;
        type.insert(p[i].second);
      }
    }
    chmax(ans, res + type.size() * type.size());
  }
  cout << ans << '\n';
} /*
 overflow checked?
 corner case checked?
 boundary checked?
 not TLE in worst case checked?
 */