Problem p03148 (C++)

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

using namespace std;
typedef long long ll;
#define REP(i, n) for (int i = 0; i < int(n); i++)
#define FOR(i, m, n) for (int i = int(m); i < int(n); i++)
#define ALL(obj) (obj).begin(), (obj).end()
#define VI vector<int>
#define VLL vector<long long>
#define VVI vector<vector<int>>
#define VVLL vector<vector<long long>>
#define VC vector<char>
#define VS vector<string>
#define VVC vector<vector<char>>
#define fore(i, a) for (auto &i : a)

typedef pair<int, int> P;
template <class T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
const int INF = 2147483647;
const ll INFL = 9223372036854775807;
const ll mod = 1000000007;

int main() {
  ll n, k;
  cin >> n >> k;
  VVLL v(n, VLL(2));
  REP(i, n) cin >> v[i][1] >> v[i][0];
  REP(i, n) v[i][1]--;

  sort(ALL(v));
  reverse(ALL(v));

  VLL tmp;
  VLL a(n, 0);
  ll ans = 0;
  ll res = 0;
  ll cnt = 0;

  REP(i, k) {
    if (a[v[i][1]] == 0) {
      a[v[i][1]] = v[i][0];
      cnt++;
      res += v[i][0];
    } else {
      res += v[i][0];
      tmp.push_back(min(a[v[i][1]], v[i][0]));
      a[v[i][1]] = max(a[v[i][1]], v[i][0]);
    }
  }

  ans = res + cnt * cnt;

  reverse(ALL(tmp));

  ll l = tmp.size();

  ll b = 0;
  FOR(i, k, n) {
    if (a[v[i][1]])
      continue;
    res -= tmp[b];
    res += v[i][0];
    a[v[i][1]] = 1;
    cnt++;
    b++;
    ans = max(ans, res + cnt * cnt);
    if (b == l)
      break;
  }

  cout << ans << endl;

  return 0;
}

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

using namespace std;
typedef long long ll;
#define REP(i, n) for (int i = 0; i < int(n); i++)
#define FOR(i, m, n) for (int i = int(m); i < int(n); i++)
#define ALL(obj) (obj).begin(), (obj).end()
#define VI vector<int>
#define VLL vector<long long>
#define VVI vector<vector<int>>
#define VVLL vector<vector<long long>>
#define VC vector<char>
#define VS vector<string>
#define VVC vector<vector<char>>
#define fore(i, a) for (auto &i : a)

typedef pair<int, int> P;
template <class T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
const int INF = 2147483647;
const ll INFL = 9223372036854775807;
const ll mod = 1000000007;

int main() {
  ll n, k;
  cin >> n >> k;
  VVLL v(n, VLL(2));
  REP(i, n) cin >> v[i][1] >> v[i][0];
  REP(i, n) v[i][1]--;

  sort(ALL(v));
  reverse(ALL(v));

  VLL tmp;
  VLL a(n, 0);
  ll ans = 0;
  ll res = 0;
  ll cnt = 0;

  REP(i, k) {
    if (a[v[i][1]] == 0) {
      a[v[i][1]] = v[i][0];
      cnt++;
      res += v[i][0];
    } else {
      res += v[i][0];
      tmp.push_back(min(a[v[i][1]], v[i][0]));
      a[v[i][1]] = max(a[v[i][1]], v[i][0]);
    }
  }

  ans = res + cnt * cnt;

  reverse(ALL(tmp));

  ll l = tmp.size();

  ll b = 0;
  if (l == 0) {
    cout << ans << endl;
    return 0;
  }
  FOR(i, k, n) {
    if (a[v[i][1]])
      continue;
    res -= tmp[b];
    res += v[i][0];
    a[v[i][1]] = 1;
    cnt++;
    b++;
    ans = max(ans, res + cnt * cnt);
    if (b == l)
      break;
  }

  cout << ans << endl;

  return 0;
}