Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
#define FOR(i, m, n) for (int i = (m); i < (n); i++)
#define FORR(i, m, n) for (int i = (m); i >= (n); i--)
#define REP(i, n) FOR(i, 0, (n))
#define REPR(i, n) FORR(i, (n)-1, 0)
#define REP1(i, n) FOR(i, 1, (n) + 1)
#define ALL(c) (c).begin(), (c).end()
#define sz(v) (int)v.size()
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
const int MOD = 1000000007;
const int INF = 1000000001;
const ll LINF = 1000000001000000001LL;
const char EOL = '\n';

void solve();

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cout << fixed << setprecision(numeric_limits<double>::max_digits10);
  solve();
  return 0;
}

void solve() {
  int n, k;
  cin >> n >> k;
  vector<pair<ll, int>> s(n);
  int t;
  ll d;
  REP(i, n) {
    cin >> t >> d;
    s[i] = make_pair(d, t);
  }
  sort(ALL(s));
  reverse(ALL(s));
  map<int, ll> v;
  deque<ll> q;
  ll tmp = 0;
  REP(i, k) {
    if (v.count(s[i].second) == 0) {
      v[s[i].second] = s[i].first;
    } else {
      q.push_front(s[i].first);
    }
    tmp += s[i].first;
  }
  ll x = sz(v);
  tmp += x * x;
  ll ans = tmp;
  FOR(i, k, n) {
    if (v.count(s[i].second) != 0)
      continue;
    ll r = q.front();
    tmp = tmp - r - (x * x) + s[i].first + ((x + 1) * (x + 1));
    q.pop_front();
    v[s[i].second] = s[i].first;
    x++;
    chmax(ans, tmp);
  }
  cout << ans << EOL;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

using ll = long long;
#define FOR(i, m, n) for (int i = (m); i < (n); i++)
#define FORR(i, m, n) for (int i = (m); i >= (n); i--)
#define REP(i, n) FOR(i, 0, (n))
#define REPR(i, n) FORR(i, (n)-1, 0)
#define REP1(i, n) FOR(i, 1, (n) + 1)
#define ALL(c) (c).begin(), (c).end()
#define sz(v) (int)v.size()
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
const int MOD = 1000000007;
const int INF = 1000000001;
const ll LINF = 1000000001000000001LL;
const char EOL = '\n';

void solve();

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  cout << fixed << setprecision(numeric_limits<double>::max_digits10);
  solve();
  return 0;
}

void solve() {
  int n, k;
  cin >> n >> k;
  vector<pair<ll, int>> s(n);
  int t;
  ll d;
  REP(i, n) {
    cin >> t >> d;
    s[i] = make_pair(d, t);
  }
  sort(ALL(s));
  reverse(ALL(s));
  map<int, ll> v;
  deque<ll> q;
  ll tmp = 0;
  REP(i, k) {
    if (v.count(s[i].second) == 0) {
      v[s[i].second] = s[i].first;
    } else {
      q.push_front(s[i].first);
    }
    tmp += s[i].first;
  }
  ll x = sz(v);
  tmp += x * x;
  ll ans = tmp;
  FOR(i, k, n) {
    if (q.empty())
      break;
    if (v.count(s[i].second) != 0)
      continue;
    ll r = q.front();
    tmp = tmp - r - (x * x) + s[i].first + ((x + 1) * (x + 1));
    q.pop_front();
    v[s[i].second] = s[i].first;
    x++;
    chmax(ans, tmp);
  }
  cout << ans << EOL;
}