Problem p03148 (C++)

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

#define rep(i, n) for (int i = 0; i < int(n); i++)
#define rrep(i, n) for (int i = int(n) - 1; i >= 0; i--)
#define reps(i, n) for (int i = 1; i <= int(n); i++)
#define rreps(i, n) for (int i = int(n); i >= 1; i--)
#define repi(i, a, b) for (int i = (a); i < int(b); i++)
#define all(a) (a).begin(), (a).end()
#define bit(b) (1ll << (b))
#define uniq(v) (v).erase(unique(all(v)), (v).end())

using namespace std;
using i32 = int;
using i64 = long long;
using f80 = long double;
using vi32 = vector<i32>;
using vi64 = vector<i64>;
using vf80 = vector<f80>;
using vstr = vector<string>;

void ioinit() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout << fixed << setprecision(16);
}
i64 gcd(i64 a, i64 b) {
  if (a % b == 0)
    return b;
  else
    return gcd(b, a % b);
}
i64 lcm(i64 a, i64 b) { return a / gcd(a, b) * b; }
template <class T> bool amax(T &x, T y) {
  if (x < y) {
    x = y;
    return 1;
  } else
    return 0;
}
template <class T> bool amin(T &x, T y) {
  if (x > y) {
    x = y;
    return 1;
  } else
    return 0;
}
template <class T> i64 lbi(vector<T> &v, T e) {
  return lower_bound(all(v), e) - v.begin();
}
template <class T> i64 ubi(vector<T> &v, T e) {
  return upper_bound(all(v), e) - v.begin();
}

int main() {
  ioinit();
  int n, k;
  cin >> n >> k;
  vector<int> v0, v1;
  {
    vector<vector<int>> v(100010);
    rep(i, n) {
      int t, d;
      cin >> t >> d;
      v[t].push_back(d);
    }
    rep(i, 100001) {
      if (v[i].size() == 0)
        continue;
      sort(all(v[i]));
      v1.push_back(v[i][v[i].size() - 1]);
      rep(j, v[i].size() - 1) v0.push_back(v[i][j]);
    }
  }
  sort(all(v0));
  reverse(all(v0));
  sort(all(v1));
  reverse(all(v1));
  {
    i64 ans = 0, sum = 0;
    i32 cnt = 0, kind = 0;
    rep(i, min(k, (int)v0.size())) {
      sum += v0[i];
      cnt++;
    }
    rep(i, max(0, k - cnt)) {
      sum += v1[i];
      kind++;
    }
    amax(ans, sum + kind * kind);
    while (cnt > 0) {
      sum -= v0[--cnt];
      sum += v1[kind++];
      amax(ans, sum + (i64)kind * kind);
    }
    cout << ans << endl;
  }
  return 0;
}


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

#define rep(i, n) for (int i = 0; i < int(n); i++)
#define rrep(i, n) for (int i = int(n) - 1; i >= 0; i--)
#define reps(i, n) for (int i = 1; i <= int(n); i++)
#define rreps(i, n) for (int i = int(n); i >= 1; i--)
#define repi(i, a, b) for (int i = (a); i < int(b); i++)
#define all(a) (a).begin(), (a).end()
#define bit(b) (1ll << (b))
#define uniq(v) (v).erase(unique(all(v)), (v).end())

using namespace std;
using i32 = int;
using i64 = long long;
using f80 = long double;
using vi32 = vector<i32>;
using vi64 = vector<i64>;
using vf80 = vector<f80>;
using vstr = vector<string>;

void ioinit() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout << fixed << setprecision(16);
}
i64 gcd(i64 a, i64 b) {
  if (a % b == 0)
    return b;
  else
    return gcd(b, a % b);
}
i64 lcm(i64 a, i64 b) { return a / gcd(a, b) * b; }
template <class T> bool amax(T &x, T y) {
  if (x < y) {
    x = y;
    return 1;
  } else
    return 0;
}
template <class T> bool amin(T &x, T y) {
  if (x > y) {
    x = y;
    return 1;
  } else
    return 0;
}
template <class T> i64 lbi(vector<T> &v, T e) {
  return lower_bound(all(v), e) - v.begin();
}
template <class T> i64 ubi(vector<T> &v, T e) {
  return upper_bound(all(v), e) - v.begin();
}

int main() {
  ioinit();
  int n, k;
  cin >> n >> k;
  vector<int> v0, v1;
  {
    vector<vector<int>> v(100010);
    rep(i, n) {
      int t, d;
      cin >> t >> d;
      v[t].push_back(d);
    }
    rep(i, 100001) {
      if (v[i].size() == 0)
        continue;
      sort(all(v[i]));
      v1.push_back(v[i][v[i].size() - 1]);
      rep(j, v[i].size() - 1) v0.push_back(v[i][j]);
    }
  }
  sort(all(v0));
  reverse(all(v0));
  sort(all(v1));
  reverse(all(v1));
  {
    i64 ans = 0, sum = 0;
    i32 cnt = 0, kind = 0;
    rep(i, min(k, (int)v0.size())) {
      sum += v0[i];
      cnt++;
    }
    rep(i, max(0, k - cnt)) {
      sum += v1[i];
      kind++;
    }
    amax(ans, sum + (i64)kind * kind);
    while (cnt > 0 && kind < v1.size()) {
      sum -= v0[--cnt];
      sum += v1[kind++];
      amax(ans, sum + (i64)kind * kind);
    }
    cout << ans << endl;
  }
  return 0;
}