prithivMLmods/Coder-Stat huggingface dataset:
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <functional>
#include <iostream>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<ll, int> P;
const int MAXN = 1e5 + 10;

int k, n;
P sushi[MAXN];
vector<int> removable;
vector<ll> v[MAXN];
vector<ll> sumV;

void solve() {
  sort(sushi, sushi + n, greater<P>());
  ll sum = 0;
  ll typeNum = 0;
  for (int i = 0; i < k; i++) {
    sum += sushi[i].first;
    if (v[sushi[i].second].size() > 0) {
      removable.push_back(sushi[i].second);
    } else {
      typeNum++;
    }
    v[sushi[i].second].push_back(sushi[i].first);
  }
  sumV.push_back(sum + typeNum * typeNum);
  for (int i = k; i < n; i++) {
    if (v[sushi[i].second].size() == 0) {
      int removeType = removable.back();
      sum -= v[removeType].back();
      v[removeType].pop_back();
      removable.pop_back();
      typeNum++;
      sum += sushi[i].first;
      v[sushi[i].second].push_back(sushi[i].first);
      sumV.push_back(sum + typeNum * typeNum);
    }
  }
  sort(sumV.begin(), sumV.end(), greater<ll>());
  cout << sumV.front() << endl;
}

int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    cin >> sushi[i].second >> sushi[i].first;
  }
  solve();
}


Fixed submission:
#include <algorithm>
#include <functional>
#include <iostream>
#include <vector>

using namespace std;
typedef long long ll;
typedef pair<ll, int> P;
const int MAXN = 1e5 + 10;

int k, n;
P sushi[MAXN];
vector<int> removable;
vector<ll> v[MAXN];
vector<ll> sumV;

void solve() {
  sort(sushi, sushi + n, greater<P>());
  ll sum = 0;
  ll typeNum = 0;
  for (int i = 0; i < k; i++) {
    sum += sushi[i].first;
    if (v[sushi[i].second].size() > 0) {
      removable.push_back(sushi[i].second);
    } else {
      typeNum++;
    }
    v[sushi[i].second].push_back(sushi[i].first);
  }
  sumV.push_back(sum + typeNum * typeNum);
  for (int i = k; i < n && !removable.empty(); i++) {
    if (v[sushi[i].second].size() == 0) {
      int removeType = removable.back();
      sum -= v[removeType].back();
      v[removeType].pop_back();
      removable.pop_back();
      typeNum++;
      sum += sushi[i].first;
      v[sushi[i].second].push_back(sushi[i].first);
      sumV.push_back(sum + typeNum * typeNum);
    }
  }
  sort(sumV.begin(), sumV.end(), greater<ll>());
  cout << sumV.front() << endl;
}

int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    cin >> sushi[i].second >> sushi[i].first;
  }
  solve();
}
Problem p03148 (C++)

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

int main() {
  int64_t N, K, t, d;
  cin >> N >> K;
  vector<priority_queue<int64_t>> vpq(MAXN);
  for (int i = 0; i < N; i++) {
    cin >> t >> d;
    vpq[t].push(d);
  }
  priority_queue<pair<int64_t, int64_t>> pq;
  for (int i = 0; i < MAXN; i++) {
    if (vpq[i].empty())
      continue;
    pq.emplace(vpq[i].top(), 1);
    vpq[i].pop();
    while (vpq[i].size()) {
      pq.emplace(vpq[i].top(), 0);
      vpq[i].pop();
    }
  }
  int64_t tmp = 0, ans = 0;
  priority_queue<int64_t, vector<int64_t>, greater<int64_t>> pq0, pq1;
  for (int i = 0; i < K; i++) {
    tmp += pq.top().first;
    if (pq.top().second)
      pq1.push(pq.top().first);
    else
      pq0.push(pq.top().first);
    pq.pop();
  }
  tmp += pq1.size() * pq1.size();
  ans = max(ans, tmp);
  while (pq.size()) {
    if (pq.top().second) {
      tmp -= pq1.size() * pq1.size();
      tmp -= pq0.top();
      pq0.pop();
      tmp += pq.top().first;
      pq1.push(pq.top().first);
      tmp += pq1.size() * pq1.size();
    }
    ans = max(ans, tmp);
    pq.pop();
  }
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100010;

int main() {
  int64_t N, K, t, d;
  cin >> N >> K;
  vector<priority_queue<int64_t>> vpq(MAXN);
  for (int i = 0; i < N; i++) {
    cin >> t >> d;
    vpq[t].push(d);
  }
  priority_queue<pair<int64_t, int64_t>> pq;
  for (int i = 0; i < MAXN; i++) {
    if (vpq[i].empty())
      continue;
    pq.emplace(vpq[i].top(), 1);
    vpq[i].pop();
    while (vpq[i].size()) {
      pq.emplace(vpq[i].top(), 0);
      vpq[i].pop();
    }
  }
  int64_t tmp = 0, ans = 0;
  priority_queue<int64_t, vector<int64_t>, greater<int64_t>> pq0, pq1;
  for (int i = 0; i < K; i++) {
    tmp += pq.top().first;
    if (pq.top().second)
      pq1.push(pq.top().first);
    else
      pq0.push(pq.top().first);
    pq.pop();
  }
  tmp += pq1.size() * pq1.size();
  ans = max(ans, tmp);
  while (pq.size()) {
    if (pq.top().second) {
      tmp -= pq1.size() * pq1.size();
      if (pq0.empty()) {
        cout << ans << endl;
        return 0;
      }
      tmp -= pq0.top();
      pq0.pop();
      tmp += pq.top().first;
      pq1.push(pq.top().first);
      tmp += pq1.size() * pq1.size();
    }
    ans = max(ans, tmp);
    pq.pop();
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
using uint = unsigned int;
using lint = long long int;
using ulint = unsigned long long int;
template <class T = int> using V = vector<T>;
template <class T = int> using VV = V<V<T>>;
template <class T, class U> void assign(V<T> &v, int n, const U &a) {
  v.assign(n, a);
}
template <class T, class... Args>
void assign(V<T> &v, int n, const Args &...args) {
  v.resize(n);
  for (auto &&e : v)
    assign(e, args...);
}

int main() {
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  V<> t(n), d(n);
  for (int i = 0; i < n; ++i)
    cin >> t[i] >> d[i], --t[i];
  V<> mx(n);
  for (int i = 0; i < n; ++i) {
    mx[t[i]] = max(mx[t[i]], d[i]);
  }
  auto _mx = mx;
  sort(rbegin(_mx), rend(_mx));
  V<lint> cmx(n + 1);
  for (int i = 0; i < n; ++i) {
    cmx[i + 1] = cmx[i] + _mx[i];
  }
  for (int i = 0; i < n; ++i) {
    if (mx[t[i]] == d[i]) {
      mx[t[i]] = d[i] = 0;
    }
  }
  sort(rbegin(d), rend(d));
  V<lint> cd(n + 1);
  for (int i = 0; i < n; ++i) {
    cd[i + 1] = cd[i] + d[i];
  }
  lint res = 0;
  int m = set<int>(begin(t), end(t)).size();
  for (lint x = 1; x <= m; ++x) {
    res = max(res, cmx[x] + cd[k - x] + x * x);
  }
  cout << res << '\n';
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
using uint = unsigned int;
using lint = long long int;
using ulint = unsigned long long int;
template <class T = int> using V = vector<T>;
template <class T = int> using VV = V<V<T>>;
template <class T, class U> void assign(V<T> &v, int n, const U &a) {
  v.assign(n, a);
}
template <class T, class... Args>
void assign(V<T> &v, int n, const Args &...args) {
  v.resize(n);
  for (auto &&e : v)
    assign(e, args...);
}

int main() {
  cin.tie(nullptr);
  ios_base::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  V<> t(n), d(n);
  for (int i = 0; i < n; ++i)
    cin >> t[i] >> d[i], --t[i];
  V<> mx(n);
  for (int i = 0; i < n; ++i) {
    mx[t[i]] = max(mx[t[i]], d[i]);
  }
  auto _mx = mx;
  sort(rbegin(_mx), rend(_mx));
  V<lint> cmx(n + 1);
  for (int i = 0; i < n; ++i) {
    cmx[i + 1] = cmx[i] + _mx[i];
  }
  for (int i = 0; i < n; ++i) {
    if (mx[t[i]] == d[i]) {
      mx[t[i]] = d[i] = 0;
    }
  }
  sort(rbegin(d), rend(d));
  V<lint> cd(n + 1);
  for (int i = 0; i < n; ++i) {
    cd[i + 1] = cd[i] + d[i];
  }
  lint res = 0;
  int m = set<int>(begin(t), end(t)).size();
  for (lint x = 1; x <= min(m, k); ++x) {
    res = max(res, cmx[x] + cd[k - x] + x * x);
  }
  cout << res << '\n';
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <string>
#include <vector>

using namespace std;
typedef long long ll;
const int MOD = 1000000007;
#define MAX_N (100010)

int main() {
  int N, K;
  scanf("%d%d", &N, &K);

  map<int, vector<int>> sushi;
  for (int i = 0; i < N; i++) {
    int t, d;
    scanf("%d%d", &t, &d);
    sushi[t].push_back(d);
  }

  vector<int> top, other;
  for (auto iter = sushi.begin(); iter != sushi.end(); iter++) {
    vector<int> elements = iter->second;
    sort(elements.begin(), elements.end(), greater<int>());
    top.push_back(elements[0]);
    for (int i = 1; i < (int)elements.size(); i++)
      other.push_back(elements[i]);
  }
  sort(top.begin(), top.end(), greater<int>());
  sort(other.begin(), other.end(), greater<int>());

  ll res = 0;
  for (int k = 1; k <= K; k++) {
    if (k > (int)top.size())
      continue;

    ll local = 0;
    for (int i = 0; i < k; i++)
      local += (ll)top[i];
    for (int i = k; i < K; i++)
      local += (ll)other[i - k];
    local += (ll)k * (ll)k;

    res = max(res, local);
  }

  printf("%lld\n", res);

  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <string>
#include <vector>

using namespace std;
typedef long long ll;
const int MOD = 1000000007;
#define MAX_N (100010)

int main() {
  int N, K;
  scanf("%d%d", &N, &K);

  map<int, vector<int>> sushi;
  for (int i = 0; i < N; i++) {
    int t, d;
    scanf("%d%d", &t, &d);
    sushi[t].push_back(d);
  }

  vector<int> top, other;
  for (auto iter = sushi.begin(); iter != sushi.end(); iter++) {
    vector<int> elements = iter->second;
    sort(elements.begin(), elements.end(), greater<int>());
    top.push_back(elements[0]);
    for (int i = 1; i < (int)elements.size(); i++)
      other.push_back(elements[i]);
  }
  sort(top.begin(), top.end(), greater<int>());
  sort(other.begin(), other.end(), greater<int>());

  ll res = 0;
  for (int k = 1; k <= K; k++) {
    if (k > (int)top.size() || K - k > (int)other.size())
      continue;

    ll local = 0;
    for (int i = 0; i < k; i++)
      local += (ll)top[i];
    for (int i = k; i < K; i++)
      local += (ll)other[i - k];
    local += (ll)k * (ll)k;

    res = max(res, local);
  }

  printf("%lld\n", res);

  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):

#include <algorithm>
#include <climits>
#include <cmath>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

const int maxn = 100002;

int n, k;
long long ans;

struct cake {
  int kind;
  long long base;
};

bool mycmp(const cake &a, const cake &b) { return a.base > b.base; }

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  cin >> n >> k;
  vector<cake> data(n);
  vector<int> freq(n + 2);

  for (int i = 0; i < n; ++i)
    cin >> data[i].kind >> data[i].base;

  sort(data.begin(), data.end(), mycmp);
  priority_queue<int, vector<int>, greater<int>> pq;

  long long tot = 0;
  for (int i = 0; i < k; ++i) {
    ans += data[i].base;
    if (!freq[data[i].kind])
      ++tot;

    ++freq[data[i].kind];

    if (freq[data[i].kind] > 1)
      pq.push(data[i].base);
  }

  long long sum = ans;

  ans += tot * tot;

  for (int i = k; i < n; ++i) {
    if (freq[data[i].kind] > 0)
      continue;

    sum += data[i].base;
    sum -= pq.top();
    pq.pop();

    ++freq[data[i].kind];
    ++tot;

    ans = max(ans, sum + tot * tot);
  }

  cout << ans << endl;

  return 0;
}

Fixed submission:

#include <algorithm>
#include <climits>
#include <cmath>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

const int maxn = 100002;

int n, k;
long long ans;

struct cake {
  int kind;
  long long base;
};

bool mycmp(const cake &a, const cake &b) { return a.base > b.base; }

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  cin >> n >> k;
  vector<cake> data(n);
  vector<int> freq(n + 2);

  for (int i = 0; i < n; ++i)
    cin >> data[i].kind >> data[i].base;

  sort(data.begin(), data.end(), mycmp);
  priority_queue<int, vector<int>, greater<int>> pq;

  long long tot = 0;
  for (int i = 0; i < k; ++i) {
    ans += data[i].base;
    if (!freq[data[i].kind])
      ++tot;

    ++freq[data[i].kind];

    if (freq[data[i].kind] > 1)
      pq.push(data[i].base);
  }

  long long sum = ans;

  ans += tot * tot;

  for (int i = k; i < n && !pq.empty(); ++i) {
    if (freq[data[i].kind] > 0)
      continue;

    sum += data[i].base;
    sum -= pq.top();
    pq.pop();

    ++freq[data[i].kind];
    ++tot;

    ans = max(ans, sum + tot * tot);
  }

  cout << ans << endl;

  return 0;
}
Problem p03148 (C++)

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

using ll = long long int;

const int MAX = (int)(1e5 + 5);
const ll INF = (ll)(1e10 + 5);

int n, k;

vector<pair<ll, int>> td;

map<int, int> mp;

priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<pair<ll, int>>>
    pque;

vector<ll> ans;

int main(void) {
  // Here your code !

  scanf("%d %d", &n, &k);

  for (int i = 0; i < n; ++i) {
    int t;
    ll d;
    scanf("%d %lld", &t, &d);

    td.push_back(make_pair(d, t));
  }

  sort(td.begin(), td.end(), greater<pair<ll, int>>());

  ll tmp_sum = 0;

  for (int i = 0; i < k; ++i) {
    mp[td[i].second] += 1;
    if (mp[td[i].second] > 1) {
      pque.push(make_pair(td[i].first, td[i].second));
    }
    tmp_sum += td[i].first;
  }

  ans.push_back(tmp_sum + pow(mp.size(), 2));

  for (int i = k; i < n; ++i) {
    if (mp.find(td[i].second) == mp.end()) {
      pair<ll, int> e = pque.top();
      pque.pop();

      mp[e.second] -= 1;
      tmp_sum -= e.first;

      mp[td[i].second] += 1;
      tmp_sum += td[i].first;

      ans.push_back(tmp_sum + powl(mp.size(), 2));
    }
  }

  ll last_ans = *max_element(ans.begin(), ans.end());

  printf("%lld\n", last_ans);

  return 0;
}


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

using ll = long long int;

const int MAX = (int)(1e5 + 5);
const ll INF = (ll)(1e10 + 5);

int n, k;

vector<pair<ll, int>> td;

map<int, int> mp;

priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<pair<ll, int>>>
    pque;

vector<ll> ans;

int main(void) {
  // Here your code !

  scanf("%d %d", &n, &k);

  for (int i = 0; i < n; ++i) {
    int t;
    ll d;
    scanf("%d %lld", &t, &d);

    td.push_back(make_pair(d, t));
  }

  sort(td.begin(), td.end(), greater<pair<ll, int>>());

  ll tmp_sum = 0;

  for (int i = 0; i < k; ++i) {
    mp[td[i].second] += 1;
    if (mp[td[i].second] > 1) {
      pque.push(make_pair(td[i].first, td[i].second));
    }
    tmp_sum += td[i].first;
  }

  ans.push_back(tmp_sum + pow(mp.size(), 2));

  for (int i = k; i < n; ++i) {
    if (mp.find(td[i].second) == mp.end() && !pque.empty()) {
      pair<ll, int> e = pque.top();
      pque.pop();

      mp[e.second] -= 1;
      tmp_sum -= e.first;

      mp[td[i].second] += 1;
      tmp_sum += td[i].first;

      ans.push_back(tmp_sum + powl(mp.size(), 2));
    }
  }

  ll last_ans = *max_element(ans.begin(), ans.end());

  printf("%lld\n", last_ans);

  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <vector>

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(a) (a).begin(), (a).end()

typedef long long lint;

using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> v[N];
  REP(i, N) {
    int t, d;
    cin >> t >> d;
    t--;
    v[t].push_back(d);
  }
  vector<lint> a, b;
  REP(i, N) {
    if (v[i].empty())
      continue;
    sort(ALL(v[i]), greater<int>());
    a.push_back(v[i][0]);
    FOR(j, 1, v[i].size()) b.push_back(v[i][j]);
  }
  sort(ALL(a), greater<lint>());
  sort(ALL(b), greater<lint>());
  // for(auto i:a)cerr<<i<<" ";
  // cerr<<endl;
  // for(auto i:b)cerr<<i<<" ";
  // cerr<<endl;
  vector<lint> sa, sb;
  sa.push_back(0);
  sb.push_back(0);
  for (auto i : a)
    sa.push_back(i + sa.back());
  for (auto i : b)
    sb.push_back(i + sb.back());
  // for(auto i:sa)cerr<<i<<" ";
  // cerr<<endl;
  // for(auto i:sb)cerr<<i<<" ";
  // cerr<<endl;
  lint ans = 0;
  FOR(i, 1, K + 1) {
    lint temp = i * i;
    temp += sa[i];
    temp += sb[K - i];
    ans = max(ans, temp);
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string.h>
#include <vector>

#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define REP(i, n) FOR(i, 0, n)
#define ALL(a) (a).begin(), (a).end()

typedef long long lint;

using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> v[N];
  REP(i, N) {
    int t, d;
    cin >> t >> d;
    t--;
    v[t].push_back(d);
  }
  vector<lint> a, b;
  REP(i, N) {
    if (v[i].empty())
      continue;
    sort(ALL(v[i]), greater<int>());
    a.push_back(v[i][0]);
    FOR(j, 1, v[i].size()) b.push_back(v[i][j]);
  }
  sort(ALL(a), greater<lint>());
  sort(ALL(b), greater<lint>());
  // for(auto i:a)cerr<<i<<" ";
  // cerr<<endl;
  // for(auto i:b)cerr<<i<<" ";
  // cerr<<endl;
  vector<lint> sa, sb;
  sa.push_back(0);
  sb.push_back(0);
  for (auto i : a)
    sa.push_back(i + sa.back());
  for (auto i : b)
    sb.push_back(i + sb.back());
  // for(auto i:sa)cerr<<i<<" ";
  // cerr<<endl;
  // for(auto i:sb)cerr<<i<<" ";
  // cerr<<endl;
  lint ans = 0;
  FOR(i, 1, K + 1) {
    if (sa.size() < i + 1)
      continue;
    if (sb.size() < K - i + 1)
      continue;
    lint temp = (lint)i * i;
    temp += sa[i];
    temp += sb[K - i];
    ans = max(ans, temp);
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
#define sz(x) (int)x.size()

int main() {
  int n;
  cin >> n;
  int k;
  cin >> k;
  vector<pair<int, int>> td(n);
  for (int i = 0; i < n; i++) {
    cin >> td[i].second >> td[i].first;
  }
  sort(td.rbegin(), td.rend());
  //~ for (int i = 0; i < n; i++) {
  //~ cout << td[i].first << " " << td[i].second << endl;
  //~ }
  //~ cout << endl;
  vector<int> un;
  vector<ll> nun;
  ll un_sum = 0;
  vector<bool> used(n + 1);
  ll ans = 0;
  nun.push_back(0);
  for (int i = 0; i < n; i++) {
    if (used[td[i].second]) {
      if (nun.empty())
        nun.push_back(td[i].first);
      else
        nun.push_back(td[i].first + nun.back());
    } else {
      used[td[i].second] = true;
      un.push_back(td[i].first);
      un_sum += td[i].first;
    }
    int x = un.size();
    int need = k - x;
    if (sz(nun) > need) {
      ll total = (ll)x * x;
      total += un_sum;
      total += nun[need];
      ans = max(ans, total);
    }
  }
  cout << ans;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
typedef long long ll;
using namespace std;
#define sz(x) (int)x.size()

int main() {
  int n;
  cin >> n;
  int k;
  cin >> k;
  vector<pair<int, int>> td(n);
  for (int i = 0; i < n; i++) {
    cin >> td[i].second >> td[i].first;
  }
  sort(td.rbegin(), td.rend());
  //~ for (int i = 0; i < n; i++) {
  //~ cout << td[i].first << " " << td[i].second << endl;
  //~ }
  //~ cout << endl;
  vector<int> un;
  vector<ll> nun;
  ll un_sum = 0;
  vector<bool> used(n + 1);
  ll ans = 0;
  nun.push_back(0);
  for (int i = 0; i < n; i++) {
    if (used[td[i].second]) {
      if (nun.empty())
        nun.push_back(td[i].first);
      else
        nun.push_back(td[i].first + nun.back());
    } else {
      used[td[i].second] = true;
      un.push_back(td[i].first);
      un_sum += td[i].first;
    }
    int x = un.size();
    int need = k - x;
    if (sz(nun) > need && need >= 0) {
      ll total = (ll)x * x;
      total += un_sum;
      total += nun[need];
      ans = max(ans, total);
    }
  }
  cout << ans;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
priority_queue<ll, vector<ll>, greater<ll>> que;
vector<P> vec;
ll n, k, t, d, sum, shu, a[100010], ans;
bool jud[100010];
int main(void) {
  // Your code here!
  cin >> n >> k;
  for (ll i = 0; i < n; i++) {
    cin >> t >> d;
    vec.push_back(P(d, t));
  }
  sort(vec.begin(), vec.end());
  reverse(vec.begin(), vec.end());
  for (ll i = 0; i < k; i++) {
    if (jud[vec[i].second]) {
      que.push(vec[i].first);
    } else {
      jud[vec[i].second] = true;
      shu++;
    }
    sum += vec[i].first;
  }
  a[shu] = sum;
  for (ll i = k; i < n; i++) {
    if (jud[vec[i].second])
      continue;
    jud[vec[i].second] = true;
    shu++;
    a[shu] = a[shu - 1] + vec[i].first - que.top();
    que.pop();
  }
  for (ll i = 0; i <= n; i++) {
    if (a[i] == 0)
      continue;
    ans = max(ans, a[i] + i * i);
  }
  cout << ans << endl;
}


Fixed submission:
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
priority_queue<ll, vector<ll>, greater<ll>> que;
vector<P> vec;
ll n, k, t, d, sum, shu, a[100010], ans;
bool jud[100010];
int main(void) {
  // Your code here!
  cin >> n >> k;
  for (ll i = 0; i < n; i++) {
    cin >> t >> d;
    vec.push_back(P(d, t));
  }
  sort(vec.begin(), vec.end());
  reverse(vec.begin(), vec.end());
  for (ll i = 0; i < k; i++) {
    if (jud[vec[i].second]) {
      que.push(vec[i].first);
    } else {
      jud[vec[i].second] = true;
      shu++;
    }
    sum += vec[i].first;
  }
  a[shu] = sum;
  for (ll i = k; i < n; i++) {
    if (jud[vec[i].second])
      continue;
    if (que.size() == 0)
      continue;
    jud[vec[i].second] = true;
    shu++;
    a[shu] = a[shu - 1] + vec[i].first - que.top();
    que.pop();
  }
  for (ll i = 0; i <= n; i++) {
    if (a[i] == 0)
      continue;
    ans = max(ans, a[i] + i * i);
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#define _USE_MATH_DEFINES

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>

using namespace std;

typedef long long ll;
typedef pair<ll, ll> P;

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

const int Maxn = 100010;

int n, k;

vector<ll> v[Maxn];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    ll t, d;
    cin >> t >> d;
    --t;
    v[t].emplace_back(d);
  }
  for (int i = 0; i < n; ++i) {
    sort(v[i].begin(), v[i].end());
  }
  vector<ll> y0, y1;
  for (int i = 0; i < n; ++i) {
    if (v[i].empty())
      continue;
    y1.emplace_back(v[i].back());
    v[i].pop_back();
    y0.insert(y0.end(), v[i].begin(), v[i].end());
  }
  sort(y0.begin(), y0.end(), greater<ll>());
  sort(y1.begin(), y1.end(), greater<ll>());

  int kind = max(0, k - (int)y0.size());
  ll sum = 0;
  for (int i = 0; i < kind; ++i) {
    sum += y1[i];
  }
  for (int i = 0; i < k - kind; ++i) {
    sum += y0[i];
  }

  ll ans = sum + 1LL * kind * kind;
  while (kind <= min(k + 1, (int)y1.size())) {
    ans = max(ans, sum + 1LL * kind * kind);
    sum += y1[kind];
    sum -= y0[k - kind - 1];
    ++kind;
  }
  cout << ans << endl;
}

Fixed submission:
#define _USE_MATH_DEFINES

#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>

using namespace std;

typedef long long ll;
typedef pair<ll, ll> P;

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

const int Maxn = 100010;

int n, k;

vector<ll> v[Maxn];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);

  cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    ll t, d;
    cin >> t >> d;
    --t;
    v[t].emplace_back(d);
  }
  for (int i = 0; i < n; ++i) {
    sort(v[i].begin(), v[i].end());
  }
  vector<ll> y0, y1;
  for (int i = 0; i < n; ++i) {
    if (v[i].empty())
      continue;
    y1.emplace_back(v[i].back());
    v[i].pop_back();
    y0.insert(y0.end(), v[i].begin(), v[i].end());
  }
  sort(y0.begin(), y0.end(), greater<ll>());
  sort(y1.begin(), y1.end(), greater<ll>());

  int kind = max(0, k - (int)y0.size());
  ll sum = 0;
  for (int i = 0; i < kind; ++i) {
    sum += y1[i];
  }
  for (int i = 0; i < k - kind; ++i) {
    sum += y0[i];
  }

  ll ans = sum + 1LL * kind * kind;
  while (kind <= min(k + 1, (int)y1.size())) {
    ans = max(ans, sum + 1LL * kind * kind);
    sum += y1[kind];
    if (kind >= k)
      break;
    sum -= y0[k - kind - 1];
    ++kind;
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < n; i++)
typedef long long ll;
int main() {
  ll n, k, t, d;
  cin >> n >> k;
  vector<pair<ll, ll>> A, B, C;
  int var[200001] = {0};
  rep(i, n) cin >> t >> d, A.push_back(make_pair(d, t));
  sort(A.begin(), A.end(), greater<pair<ll, ll>>());
  rep(i, k) B.push_back(A[i]), var[A[i].second]++;
  rep(i, n - k) C.push_back(A[i + k]);
  reverse(B.begin(), B.end());
  ll c = 0, ans = 0, sum = 0, x = 0;
  rep(i, k) sum += B[i].first;
  rep(i, n) if (var[i + 1]) x++;
  ans = sum + x * x;
  rep(i, k) {
    if (var[B[i].second] > 1) {
      while (c < n - k && var[C[c].second] != 0)
        c++;
      if (c == n - k)
        break;
      var[B[i].second]--;
      sum -= B[i].first;
      B[i].first = C[i].first;
      B[i].second = C[i].second;
      var[B[i].second]++;
      sum += B[i].first;
      x++;
    }
    ans = max(ans, sum + x * x);
  }
  cout << ans;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < n; i++)
typedef long long ll;
int main() {
  ll n, k, t, d;
  cin >> n >> k;
  vector<pair<ll, ll>> A, B, C;
  int var[200001] = {0};
  rep(i, n) cin >> t >> d, A.push_back(make_pair(d, t));
  sort(A.begin(), A.end(), greater<pair<ll, ll>>());
  rep(i, k) B.push_back(A[i]), var[A[i].second]++;
  rep(i, n - k) C.push_back(A[i + k]);
  reverse(B.begin(), B.end());
  ll c = 0, ans = 0, sum = 0, x = 0;
  rep(i, k) sum += B[i].first;
  rep(i, n) if (var[i + 1]) x++;
  ans = sum + x * x;
  rep(i, k) {
    if (var[B[i].second] > 1) {
      while (c < n - k && var[C[c].second] != 0)
        c++;
      if (c == n - k)
        break;
      var[B[i].second]--;
      sum -= B[i].first;
      B[i].first = C[c].first;
      B[i].second = C[c].second;
      var[B[i].second]++;
      sum += B[i].first;
      x++;
    }
    ans = max(ans, sum + x * x);
  }
  cout << ans;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; i++)
#define rrep(i, n) for (int i = n - 1; i >= 0; i--)
using namespace std;
#define INF ((1 << 30) - 1)
#define LINF (1LL << 60)
#define EPS (1e-10)
typedef long long ll;
typedef pair<ll, ll> P;
const int MOD = 1000000007;
const int MOD2 = 998244353;

// ref : https://atcoder.jp/contests/abc116/submissions/4327163

int main() {
  ll n, k;
  cin >> n >> k;
  vector<P> sushi(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    sushi.push_back({d, t});
  }
  sort(sushi.begin(), sushi.end());
  reverse(sushi.begin(), sushi.end());

  ll ans = 0;
  vector<ll> used(n, 0);
  stack<ll> sta;
  ll res = 0;
  ll kind = 0;
  rep(i, k) {
    res += sushi[i].first;
    if (used[sushi[i].second]++ == 0)
      kind++;
    else
      sta.push(sushi[i].first);
  }
  ans = res + kind * kind;
  for (int i = k; i < n && sta.size(); i++) {
    if (used[sushi[i].second]++)
      continue;
    ll a = sta.top();
    sta.pop();
    res -= a;
    res += sushi[i].first;
    kind++;
    ans = max(ans, res + kind * kind);
  }
  cout << ans << endl;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; i++)
#define rrep(i, n) for (int i = n - 1; i >= 0; i--)
using namespace std;
#define INF ((1 << 30) - 1)
#define LINF (1LL << 60)
#define EPS (1e-10)
typedef long long ll;
typedef pair<ll, ll> P;
const int MOD = 1000000007;
const int MOD2 = 998244353;

// ref : https://atcoder.jp/contests/abc116/submissions/4327163

int main() {
  ll n, k;
  cin >> n >> k;
  vector<P> sushi(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    t--;
    sushi[i] = {d, t};
  }
  sort(sushi.begin(), sushi.end());
  reverse(sushi.begin(), sushi.end());

  ll ans = 0;
  vector<ll> used(n, 0);
  stack<ll> sta;
  ll res = 0;
  ll kind = 0;
  rep(i, k) {
    res += sushi[i].first;
    if (used[sushi[i].second]++ == 0)
      kind++;
    else
      sta.push(sushi[i].first);
  }
  ans = res + kind * kind;
  for (int i = k; i < n && sta.size(); i++) {
    if (used[sushi[i].second]++)
      continue;
    ll a = sta.top();
    sta.pop();
    res -= a;
    res += sushi[i].first;
    kind++;
    ans = max(ans, res + kind * kind);
  }
  cout << ans << endl;

  return 0;
}
Problem p03148 (C++)

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

using namespace std;
using ll = long long;

int main(int argc, const char *argv[]) {
  int n, k;
  cin >> n >> k;

  map<int, vector<ll>> sushi;
  for (int i = 0; i < n; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    sushi[t].push_back(d);
  }

  vector<ll> g1, g2;
  for (auto &p : sushi) {
    auto &vd = p.second;
    auto max_it = max_element(vd.begin(), vd.end());
    g1.push_back(*max_it);
    for (auto it = vd.begin(); it != vd.end(); it++) {
      if (it != max_it) {
        g2.push_back(*it);
      }
    }
  }

  sort(g1.begin(), g1.end(), greater<ll>());
  sort(g2.begin(), g2.end(), greater<ll>());

  vector<ll> g1_psum{0}, g2_psum{0};
  partial_sum(g1.begin(), g1.end(), back_inserter(g1_psum));
  partial_sum(g2.begin(), g2.end(), back_inserter(g2_psum));
  ll ans = 0;
  for (ll i = 1; i <= k && i <= g1.size(); ++i) {
    ll sum = i * i + g1_psum[i] + g2_psum[k - i];
    ans = max(ans, sum);
  }

  cout << ans << '\n';

  return 0;
}

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

using namespace std;
using ll = long long;

int main(int argc, const char *argv[]) {
  int n, k;
  cin >> n >> k;

  map<int, vector<ll>> sushi;
  for (int i = 0; i < n; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    sushi[t].push_back(d);
  }

  vector<ll> g1, g2;
  for (auto &p : sushi) {
    auto &vd = p.second;
    auto max_it = max_element(vd.begin(), vd.end());
    g1.push_back(*max_it);
    for (auto it = vd.begin(); it != vd.end(); it++) {
      if (it != max_it) {
        g2.push_back(*it);
      }
    }
  }

  sort(g1.begin(), g1.end(), greater<ll>());
  sort(g2.begin(), g2.end(), greater<ll>());

  vector<ll> g1_psum{0}, g2_psum{0};
  partial_sum(g1.begin(), g1.end(), back_inserter(g1_psum));
  partial_sum(g2.begin(), g2.end(), back_inserter(g2_psum));
  ll ans = 0;
  for (ll i = 1; i <= k && i <= g1.size(); ++i) {
    if (k - i >= g2_psum.size()) {
      continue;
    }

    ll sum = i * i + g1_psum[i] + g2_psum[k - i];
    ans = max(ans, sum);
  }

  cout << ans << '\n';

  return 0;
}
Problem p03148 (C++)

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

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> in(n);
  for (int i = 0; i < n; i++) {
    cin >> in[i].first >> in[i].second;
    in[i].first--;
  }
  sort(in.begin(), in.end());
  reverse(in.begin(), in.end());
  vector<int> kind(n), to_max(n);
  for (int i = 0; i < n; i++) {
    int x = in[i].first;
    if (kind[x] == 0) {
      kind[x]++;
      to_max[i]++;
    }
  }
  vector<int> mx, els;
  for (int i = 0; i < n; i++) {
    if (to_max[i])
      mx.push_back(in[i].second);
    else
      els.push_back(in[i].second);
  }
  sort(mx.rbegin(), mx.rend());
  sort(els.rbegin(), els.rend());
  vector<long long> a((int)mx.size() + 1), b((int)els.size() + 1);
  for (int i = 0; i < (int)mx.size(); i++) {
    a[i + 1] = a[i] + mx[i];
  }
  for (int i = 0; i < (int)els.size(); i++) {
    b[i + 1] = b[i] + els[i];
  }
  long long ans = 0;
  int sa = (int)a.size() - 1;
  for (int i = 1; i <= sa; i++) {
    long long tmp = 1LL * i * i;
    if (i > k)
      break;
    tmp += a[i] + b[k - i];
    ans = max(ans, tmp);
  }
  cout << ans << endl;
  return 0;
}

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

signed main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> in(n);
  for (int i = 0; i < n; i++) {
    cin >> in[i].first >> in[i].second;
    in[i].first--;
  }
  sort(in.begin(), in.end());
  reverse(in.begin(), in.end());
  vector<int> kind(n), to_max(n);
  for (int i = 0; i < n; i++) {
    int x = in[i].first;
    if (kind[x] == 0) {
      kind[x]++;
      to_max[i]++;
    }
  }
  vector<int> mx, els;
  for (int i = 0; i < n; i++) {
    if (to_max[i])
      mx.push_back(in[i].second);
    else
      els.push_back(in[i].second);
  }
  sort(mx.rbegin(), mx.rend());
  sort(els.rbegin(), els.rend());
  vector<long long> a((int)mx.size() + 1), b((int)els.size() + 1);
  for (int i = 0; i < (int)mx.size(); i++) {
    a[i + 1] = a[i] + mx[i];
  }
  for (int i = 0; i < (int)els.size(); i++) {
    b[i + 1] = b[i] + els[i];
  }
  long long ans = 0;
  int sa = (int)a.size() - 1;
  for (int i = 1; i <= sa; i++) {
    long long tmp = 1LL * i * i;
    if (i > k)
      break;
    tmp += a[i] + b[min((int)b.size() - 1, k - i)];
    ans = max(ans, tmp);
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

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

typedef long long ll;
typedef vector<int> vi;
typedef vector<long> vl;

ll solve(int N, int K, vi &t, vi &d) {
  map<int, vi> sushi_list;
  for (int i = 0; i < N; i++) {
    sushi_list[t[i]].push_back(d[i]);
  }

  vi winner, looser;
  for (auto p : sushi_list) {
    vi q = p.second;
    sort(q.begin(), q.end());
    int w = q.back();
    winner.push_back(w);
    q.pop_back();
    looser.insert(looser.begin(), q.begin(), q.end());
  }

  sort(winner.begin(), winner.end(), greater<int>());
  sort(looser.begin(), looser.end(), greater<int>());

  int W = winner.size();
  int L = looser.size();

  vl winner_cumsum(W + 1);
  vl looser_cumsum(L + 1);
  winner_cumsum[0] = 0;
  looser_cumsum[0] = 0;

  for (int i = 0; i < W; ++i) {
    winner_cumsum[i + 1] = winner_cumsum[i] + winner[i];
  }

  for (int i = 0; i < L; ++i) {
    looser_cumsum[i + 1] = looser_cumsum[i] + looser[i];
  }

  ll best = 0;

  for (int n_unique = 1; n_unique <= W; ++n_unique) {
    if (n_unique <= K && (K - n_unique <= L)) {
      ll tmp = winner_cumsum[n_unique] + looser[K - n_unique] +
               (ll)n_unique * n_unique;
      best = max(best, tmp);
    }
  }

  return best;
}

int main() {
  int N, K;
  cin >> N >> K;
  vi t, d;
  for (int i = 0; i < N; i++) {
    int t_, d_;
    cin >> t_ >> d_;
    t.push_back(t_);
    d.push_back(d_);
  }

  ll ans = solve(N, K, t, d);
  cout << ans << endl;
  return 0;
}

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

typedef long long ll;
typedef vector<int> vi;
typedef vector<long> vl;

ll solve(int N, int K, vi &t, vi &d) {
  map<int, vi> sushi_list;
  for (int i = 0; i < N; i++) {
    sushi_list[t[i]].push_back(d[i]);
  }

  vi winner, looser;
  for (auto p : sushi_list) {
    vi q = p.second;
    sort(q.begin(), q.end());
    int w = q.back();
    winner.push_back(w);
    q.pop_back();
    looser.insert(looser.begin(), q.begin(), q.end());
  }

  sort(winner.begin(), winner.end(), greater<int>());
  sort(looser.begin(), looser.end(), greater<int>());

  int W = winner.size();
  int L = looser.size();

  vl winner_cumsum(W + 1);
  vl looser_cumsum(L + 1);
  winner_cumsum[0] = 0;
  looser_cumsum[0] = 0;

  for (int i = 0; i < W; ++i) {
    winner_cumsum[i + 1] = winner_cumsum[i] + winner[i];
  }

  for (int i = 0; i < L; ++i) {
    looser_cumsum[i + 1] = looser_cumsum[i] + looser[i];
  }

  ll best = 0;

  for (int n_unique = 1; n_unique <= W; ++n_unique) {
    if (n_unique <= K && (K - n_unique <= L)) {
      ll tmp = winner_cumsum[n_unique] + looser_cumsum[K - n_unique] +
               (ll)n_unique * n_unique;
      best = max(best, tmp);
    }
  }

  return best;
}

int main() {
  int N, K;
  cin >> N >> K;
  vi t, d;
  for (int i = 0; i < N; i++) {
    int t_, d_;
    cin >> t_ >> d_;
    t.push_back(t_);
    d.push_back(d_);
  }

  ll ans = solve(N, K, t, d);
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<int> vi;

int main() {
  int N, K;
  cin >> N >> K;
  map<int, vll> sushi_list;

  for (int i = 0; i < N; ++i) {
    int t;
    cin >> t;
    ll d;
    cin >> d;
    sushi_list[t].push_back(d);
  }

  vll winner, loser;

  for (auto p : sushi_list) {
    vll q = p.second;
    sort(q.begin(), q.end(), greater<ll>());

    int M = q.size();
    winner.push_back(q[0]);
    if (M > 1) {
      for (int i = 1; i < M; ++i)
        loser.push_back(q[i]);
    }
  }

  sort(winner.begin(), winner.end(), greater<ll>());
  sort(loser.begin(), loser.end(), greater<ll>());

  int W = winner.size(), L = loser.size();

  vll w_cum(W + 1), l_cum(L + 1);
  w_cum[0] = 0, l_cum[0] = 0;
  for (int i = 0; i < W; ++i) {
    w_cum[i + 1] = w_cum[i] + winner[i];
  }
  for (int i = 0; i < L; ++i) {
    l_cum[i + 1] = l_cum[i] + loser[i];
  }

  ll ans = 0;
  int K_ = min(K, W);
  for (int x = 1; x <= K_; ++x) {
    ll tmp = x * x + w_cum[x] + l_cum[K - x];
    ans = max(ans, tmp);
  }
  cout << ans << endl;
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<int> vi;

int main() {
  int N, K;
  cin >> N >> K;
  map<int, vll> sushi_list;

  for (int i = 0; i < N; ++i) {
    int t;
    cin >> t;
    ll d;
    cin >> d;
    sushi_list[t].push_back(d);
  }

  vll winner, loser;

  for (auto p : sushi_list) {
    vll q = p.second;
    sort(q.begin(), q.end(), greater<ll>());

    int M = q.size();
    winner.push_back(q[0]);
    if (M > 1) {
      for (int i = 1; i < M; ++i)
        loser.push_back(q[i]);
    }
  }

  sort(winner.begin(), winner.end(), greater<ll>());
  sort(loser.begin(), loser.end(), greater<ll>());

  int W = winner.size(), L = loser.size();

  vll w_cum(W + 1), l_cum(L + 1);
  w_cum[0] = 0, l_cum[0] = 0;
  for (int i = 0; i < W; ++i) {
    w_cum[i + 1] = w_cum[i] + winner[i];
  }
  for (int i = 0; i < L; ++i) {
    l_cum[i + 1] = l_cum[i] + loser[i];
  }

  ll ans = 0;
  int K_ = min(K, W);
  for (int x = 1; x <= K_; ++x) {
    if (K - x <= L) {
      ll tmp = (ll)x * x + w_cum[x] + l_cum[K - x];
      ans = max(ans, tmp);
    }
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<ll> vl;
typedef vector<vector<ll>> vvl;
typedef pair<ll, ll> P;
#define rep(i, n) for (ll i = 0; i < n; i++)
#define exrep(i, a, b) for (ll i = a; i <= b; i++)
#define out(x) cout << x << endl
#define exout(x) printf("%.10f\n", x)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define pb push_back
#define re0 return 0
const ll mod = 1000000007;
const ll INF = 1e16;
const ll MAX_N = 100010;

int main() {
  ll n, k;
  cin >> n >> k;

  vector<P> p(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    p[i] = make_pair(d, t);
  }
  sort(rall(p));

  map<ll, ll> mp;
  priority_queue<P, vector<P>, greater<P>> prQ;
  ll ans = 0;
  rep(i, k) {
    ll d = p[i].first;
    ll t = p[i].second;
    ans += d;
    mp[t]++;
    if (mp[t] >= 2) {
      prQ.push(make_pair(d, t));
    }
  }
  ll x = mp.size();
  ans += x * x;

  ll now = ans;
  exrep(i, k, n - 1) {
    ll d = p[i].first;
    ll t = p[i].second;
    if (mp[t] > 0) {
      continue;
    }
    auto z = prQ.top();
    prQ.pop();
    now += d - z.first;
    now += 2 * x + 1;
    mp[t]++;
    x++;
    chmax(ans, now);
  }

  out(ans);
  re0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <utility>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<ll> vl;
typedef vector<vector<ll>> vvl;
typedef pair<ll, ll> P;
#define rep(i, n) for (ll i = 0; i < n; i++)
#define exrep(i, a, b) for (ll i = a; i <= b; i++)
#define out(x) cout << x << endl
#define exout(x) printf("%.10f\n", x)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define pb push_back
#define re0 return 0
const ll mod = 1000000007;
const ll INF = 1e16;
const ll MAX_N = 100010;

int main() {
  ll n, k;
  cin >> n >> k;

  vector<P> p(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    p[i] = make_pair(d, t);
  }
  sort(rall(p));

  map<ll, ll> mp;
  priority_queue<P, vector<P>, greater<P>> prQ;
  ll ans = 0;
  rep(i, k) {
    ll d = p[i].first;
    ll t = p[i].second;
    ans += d;
    mp[t]++;
    if (mp[t] >= 2) {
      prQ.push(make_pair(d, t));
    }
  }
  ll x = mp.size();
  ans += x * x;

  ll now = ans;
  exrep(i, k, n - 1) {
    if (prQ.empty()) {
      break;
    }
    ll d = p[i].first;
    ll t = p[i].second;
    if (mp[t] > 0) {
      continue;
    }
    auto z = prQ.top();
    prQ.pop();
    now += d - z.first;
    now += 2 * x + 1;
    mp[t]++;
    x++;
    chmax(ans, now);
  }

  out(ans);
  re0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):

#include <bits/stdc++.h>
using namespace std;
#define inf INT_MAX
#define INF LLONG_MAX
#define ll long long
#define ull unsigned long long
#define M (int)(1e9 + 7)
#define P pair<int, int>
#define PLL pair<ll, ll>
#define FOR(i, m, n) for (int i = (int)m; i < (int)n; i++)
#define RFOR(i, m, n) for (int i = (int)m; i >= (int)n; i--)
#define rep(i, n) FOR(i, 0, n)
#define rrep(i, n) RFOR(i, n, 0)
#define all(a) a.begin(), a.end()
#define IN(a, n)                                                               \
  rep(i, n) { cin >> a[i]; }
const int vx[4] = {0, 1, 0, -1};
const int vy[4] = {1, 0, -1, 0};
#define PI 3.14159265
#define F first
#define S second
#define PB push_back
#define EB emplace_back
#define int ll
#define vi vector<int>

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

  int k, n;
  cin >> n >> k;
  priority_queue<P> q1;
  priority_queue<P, vector<P>, greater<P>> q2;

  rep(i, n) {
    int t, d;
    cin >> t >> d;
    q1.push(P(d, t));
  }

  int ans = 0, sum = 0;
  vi v(n, 0);
  int x = 0;
  rep(i, k) {
    P p = q1.top();
    q1.pop();
    sum += p.F;
    if (v[p.S] == 0)
      x++;
    v[p.S]++;
    if (v[p.S] > 1)
      q2.push(p);
  }
  ans = sum + x * x;

  rep(i, n - k) {
    P p = q1.top();
    q1.pop();
    if (v[p.S] || q2.empty())
      continue;
    P p2 = q2.top();
    q2.pop();
    sum -= p2.F;
    sum += p.F;
    v[p.S]++;
    x++;
    ans = max(ans, sum + x * x);
  }

  cout << ans << endl;
}

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define inf INT_MAX
#define INF LLONG_MAX
#define ll long long
#define ull unsigned long long
#define M (int)(1e9 + 7)
#define P pair<int, int>
#define PLL pair<ll, ll>
#define FOR(i, m, n) for (int i = (int)m; i < (int)n; i++)
#define RFOR(i, m, n) for (int i = (int)m; i >= (int)n; i--)
#define rep(i, n) FOR(i, 0, n)
#define rrep(i, n) RFOR(i, n, 0)
#define all(a) a.begin(), a.end()
#define IN(a, n)                                                               \
  rep(i, n) { cin >> a[i]; }
const int vx[4] = {0, 1, 0, -1};
const int vy[4] = {1, 0, -1, 0};
#define PI 3.14159265
#define F first
#define S second
#define PB push_back
#define EB emplace_back
#define int ll
#define vi vector<int>

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

  int k, n;
  cin >> n >> k;
  priority_queue<P> q1;
  priority_queue<P, vector<P>, greater<P>> q2;

  rep(i, n) {
    int t, d;
    cin >> t >> d;
    q1.push(P(d, t));
  }

  int ans = 0, sum = 0;
  vi v(n + 1, 0);
  int x = 0;
  rep(i, k) {
    P p = q1.top();
    q1.pop();
    sum += p.F;
    if (v[p.S] == 0)
      x++;
    v[p.S]++;
    if (v[p.S] > 1)
      q2.push(p);
  }
  ans = sum + x * x;

  rep(i, n - k) {
    P p = q1.top();
    q1.pop();
    if (v[p.S] || q2.empty())
      continue;
    P p2 = q2.top();
    q2.pop();
    sum -= p2.F;
    sum += p.F;
    v[p.S]++;
    x++;
    ans = max(ans, sum + x * x);
  }

  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

int main() {
  ll n, k, ans = 0;
  cin >> n >> k;
  vector<pair<ll, ll>> x(n);
  vector<bool> b(n, false);
  for (int i = 0; i < n; i++) {
    ll t, d;
    cin >> t >> d;
    x[i].first = d;
    x[i].second = t;
  }
  sort(x.begin(), x.end());
  reverse(x.begin(), x.end());
  map<ll, ll> mp;
  for (int i = 0; i < k; i++) {
    ans += x[i].first;
    b[i] = true;
    mp[x[i].second]++;
  }
  ans += mp.size() * mp.size();
  int j = k;
  ll sub = 0, cnt = mp.size(), sans = ans;
  for (int i = k - 1; i * (n - 1 - i) >= 0 && j < n && i >= 0;) {
    while (i >= 0 && (!b[i] || mp[x[i].second] <= 1)) {
      i--;
    }
    while (j < n && (b[j] || mp[x[j].second] >= 1)) {
      j++;
    }
    if (i * (n - 1 - i) >= 0 && j < n)
      sub = x[j].first - x[i].first + 2 * cnt + 1;
    else
      break;
    sans += sub;
    if (sans > ans)
      ans = sans;
    cnt++;
    mp[x[i].second]--;
    b[i] = false;
    mp[x[j].second]++;
    b[j] = true;
    i = j;
    j++;
  }
  cout << ans << endl;
  return 0;
}

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

int main() {
  ll n, k, ans = 0;
  cin >> n >> k;
  vector<pair<ll, ll>> x(n);
  vector<bool> b(n, false);
  for (int i = 0; i < n; i++) {
    ll t, d;
    cin >> t >> d;
    x[i].first = d;
    x[i].second = t;
  }
  sort(x.begin(), x.end());
  reverse(x.begin(), x.end());
  map<ll, ll> mp;
  for (int i = 0; i < k; i++) {
    ans += x[i].first;
    b[i] = true;
    mp[x[i].second]++;
  }
  ans += mp.size() * mp.size();
  int j = k;
  ll sub = 0, cnt = mp.size(), sans = ans;
  for (int i = k - 1; i * (n - 1 - i) >= 0 && j < n && i >= 0;) {
    while (i >= 0 && (!b[i] || mp[x[i].second] <= 1)) {
      i--;
    }
    while (j < n && (b[j] || mp[x[j].second] >= 1)) {
      j++;
    }
    if (i * (n - 1 - i) >= 0 && j < n)
      sub = x[j].first - x[i].first + 2 * cnt + 1;
    else
      break;
    sans += sub;
    if (sans > ans)
      ans = sans;
    cnt++;
    mp[x[i].second]--;
    b[i] = false;
    mp[x[j].second]++;
    b[j] = true;
    j++;
  }
  cout << ans << endl;
  return 0;
}