Problem p03148 (C++)

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

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<int, ll> pil;
typedef pair<ll, int> pli;
typedef pair<ll, ll> pll;

template <class T, class U> constexpr bool chmin(T &a, const U b) {
  if (a <= b)
    return false;
  a = b;
  return true;
}
template <class T, class U> constexpr bool chmax(T &a, const U b) {
  if (a >= b)
    return false;
  a = b;
  return true;
}
template <class T, class U> constexpr bool bit(const T n, const U k) {
  return (n >> k) & 1;
}

// デバッグ
template <class T> void Vout(vector<T> &V) {
  cout << "start" << endl;
  const int sz = V.size();
  for (int i = 0; i < sz; i++) {
    cout << i << " " << V[i] << '\n';
  }
  cout << "end" << endl;
}

//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

constexpr int MAX = 1 << 30;
constexpr ll INF = 1LL << 62;
constexpr int MOD = 1e9 + 7;

//__builtin_popcount(S);
// #define int ll
// vector<vector<int>> data(3, vector<int>(4));
// vector.resize(a,vector<int>(b,-1));
// vector<vector<vector<要素の型>>> 変数名(要素数1,
// vector<vector<要素の型>>(要素数2, vector<要素の型>(要素数3, 初期値)));

int dx[] = {1, -1, 0, 0}, dy[] = {0, 0, -1, 1};

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  int N, K;
  cin >> N >> K;
  vector<vector<ll>> V(N + 1, vector<ll>(0));
  for (int i = 0; i < N; i++) {
    int t;
    ll d;
    cin >> t >> d;
    V[t].push_back(d);
  }
  for (int i = 0; i <= N; i++) {
    sort(V[i].rbegin(), V[i].rend());
  }
  vector<ll> ne, ol;
  for (int i = 0; i <= N; i++) {
    for (int j = 0; j < (int)V[i].size(); j++) {
      if (j == 0)
        ne.push_back(V[i][j]);
      else
        ol.push_back(V[i][j]);
    }
  }
  sort(ne.rbegin(), ne.rend());
  sort(ol.rbegin(), ol.rend());

  for (int i = 1; i < (int)ne.size(); i++)
    ne[i] += ne[i - 1];
  for (int i = 1; i < (int)ol.size(); i++)
    ol[i] += ol[i - 1];
  ll ans = 0;
  for (int i = 1; i <= K; i++) {
    if (ne.size() < i)
      break;
    ll x = ne[i - 1];
    if (i != K)
      x += ol[min(K - i - 1, (int)ol.size() - 1)];
    chmax(ans, x + (ll)i * i);
  }
  cout << ans << endl;
}

Fixed submission:
#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<int, ll> pil;
typedef pair<ll, int> pli;
typedef pair<ll, ll> pll;

template <class T, class U> constexpr bool chmin(T &a, const U b) {
  if (a <= b)
    return false;
  a = b;
  return true;
}
template <class T, class U> constexpr bool chmax(T &a, const U b) {
  if (a >= b)
    return false;
  a = b;
  return true;
}
template <class T, class U> constexpr bool bit(const T n, const U k) {
  return (n >> k) & 1;
}

// デバッグ
template <class T> void Vout(vector<T> &V) {
  cout << "start" << endl;
  const int sz = V.size();
  for (int i = 0; i < sz; i++) {
    cout << i << " " << V[i] << '\n';
  }
  cout << "end" << endl;
}

//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
//+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

constexpr int MAX = 1 << 30;
constexpr ll INF = 1LL << 62;
constexpr int MOD = 1e9 + 7;

//__builtin_popcount(S);
// #define int ll
// vector<vector<int>> data(3, vector<int>(4));
// vector.resize(a,vector<int>(b,-1));
// vector<vector<vector<要素の型>>> 変数名(要素数1,
// vector<vector<要素の型>>(要素数2, vector<要素の型>(要素数3, 初期値)));

int dx[] = {1, -1, 0, 0}, dy[] = {0, 0, -1, 1};

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  int N, K;
  cin >> N >> K;
  vector<vector<ll>> V(N + 1, vector<ll>(0));
  for (int i = 0; i < N; i++) {
    int t;
    ll d;
    cin >> t >> d;
    V[t].push_back(d);
  }
  for (int i = 0; i <= N; i++) {
    sort(V[i].rbegin(), V[i].rend());
  }
  vector<ll> ne, ol;
  for (int i = 0; i <= N; i++) {
    for (int j = 0; j < (int)V[i].size(); j++) {
      if (j == 0)
        ne.push_back(V[i][j]);
      else
        ol.push_back(V[i][j]);
    }
  }
  sort(ne.rbegin(), ne.rend());
  sort(ol.rbegin(), ol.rend());

  for (int i = 1; i < (int)ne.size(); i++)
    ne[i] += ne[i - 1];
  for (int i = 1; i < (int)ol.size(); i++)
    ol[i] += ol[i - 1];
  ll ans = 0;
  for (int i = 1; i <= K; i++) {
    if (ne.size() < i)
      break;
    ll x = ne[i - 1];
    if (i != K && (int)ol.size() > 0)
      x += ol[min(K - i - 1, (int)ol.size() - 1)];
    chmax(ans, x + (ll)i * i);
  }
  cout << ans << endl;
}