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

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
using Int = long long;
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}

// INSERT ABOVE HERE
signed main() {
  Int n, k;
  cin >> n >> k;
  vector<Int> t(n), d(n);
  for (Int i = 0; i < n; i++)
    cin >> t[i] >> d[i], t[i]--;
  {
    using P = pair<Int, Int>;
    vector<P> vp;
    for (Int i = 0; i < n; i++)
      vp.emplace_back(d[i], t[i]);
    sort(vp.rbegin(), vp.rend());
    for (Int i = 0; i < n; i++)
      tie(d[i], t[i]) = vp[i];
  }

  priority_queue<Int, vector<Int>, greater<Int>> pq;

  Int res = 0;
  vector<Int> used(n, 0);
  for (Int i = 0; i < k; i++) {
    if (used[t[i]])
      pq.emplace(d[i]);
    used[t[i]] = 1;
    res += d[i];
  }
  Int cnt = 0;
  for (Int i = 0; i < n; i++)
    cnt += used[i];

  Int ans = res + cnt * cnt;
  for (Int i = k; i < n; i++) {
    if (used[t[i]])
      continue;
    used[t[i]] = 1;
    cnt++;
    res -= pq.top();
    pq.pop();
    res += d[i];
    chmax(ans, res + cnt * cnt);
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
using Int = long long;
template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
}
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
}

// INSERT ABOVE HERE
signed main() {
  Int n, k;
  cin >> n >> k;
  vector<Int> t(n), d(n);
  for (Int i = 0; i < n; i++)
    cin >> t[i] >> d[i], t[i]--;
  {
    using P = pair<Int, Int>;
    vector<P> vp;
    for (Int i = 0; i < n; i++)
      vp.emplace_back(d[i], t[i]);
    sort(vp.rbegin(), vp.rend());
    for (Int i = 0; i < n; i++)
      tie(d[i], t[i]) = vp[i];
  }

  priority_queue<Int, vector<Int>, greater<Int>> pq;

  Int res = 0;
  vector<Int> used(n, 0);
  for (Int i = 0; i < k; i++) {
    if (used[t[i]])
      pq.emplace(d[i]);
    used[t[i]] = 1;
    res += d[i];
  }
  Int cnt = 0;
  for (Int i = 0; i < n; i++)
    cnt += used[i];

  Int ans = res + cnt * cnt;
  for (Int i = k; i < n; i++) {
    if (used[t[i]])
      continue;
    if (pq.empty())
      break;
    used[t[i]] = 1;
    cnt++;
    res -= pq.top();
    pq.pop();
    res += d[i];
    chmax(ans, res + cnt * cnt);
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
n, k = map(int, input().split())
td = sorted(
    [list(map(int, input().split())) for i in range(n)],
    reverse=True,
    key=lambda x: x[1],
)

type = set()
L = []
Sum = 0
for x in td[:k]:
    Sum += x[1]
    if x[0] not in type:
        type.add(x[0])
    else:
        L.append(x[1])

L = L[::-1]
type_cnt = len(type)
ans = Sum + type_cnt**2
for x in td[k:]:
    if x[0] not in type:
        type.add(x[0])
        type_cnt += 1
        Sum = Sum - L.pop(0) + x[1]
        ans = max(ans, Sum + type_cnt**2)
print(ans)


Fixed submission:
n, k = map(int, input().split())
td = sorted(
    [list(map(int, input().split())) for i in range(n)],
    reverse=True,
    key=lambda x: x[1],
)

type = set()
L = []
Sum = 0
for x in td[:k]:
    Sum += x[1]
    if x[0] not in type:
        type.add(x[0])
    else:
        L.append(x[1])

L = L[::-1]
type_cnt = len(type)
ans = Sum + type_cnt**2
for x in td[k:]:
    if len(L) == 0:
        break
    if x[0] not in type:
        type.add(x[0])
        type_cnt += 1
        Sum = Sum - L.pop(0) + x[1]
        ans = max(ans, Sum + type_cnt**2)
print(ans)
Problem p03148 (C++)

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

using namespace std;

void solve(long long N, long long K, std::vector<long long> t,
           std::vector<long long> d) {
  vector<pair<long long, long long>> vs;
  for (int i = 0; i < K; i++) {
    vs.push_back({d[i], t[i]});
  }
  sort(vs.rbegin(), vs.rend());
  for (int i = 0; i < vs.size(); i++)
    swap(vs[i].first, vs[i].second);
  long long tot = 0;
  set<long long> ts;
  vector<long long> remCand;
  for (int i = 0; i < K; i++) {
    if (!ts.count(vs[i].first)) {
      ts.insert(vs[i].first);
    } else {
      remCand.push_back(vs[i].second);
    }
    tot += vs[i].second;
  }
  sort(remCand.rbegin(), remCand.rend());

  long long ans = tot + 1ll * ts.size() * ts.size();
  for (int i = K; i < N; i++) {
    if (!ts.count(vs[i].first) && remCand.size() > 0) {
      tot += vs[i].second;
      tot -= remCand.back();
      remCand.pop_back();
      ts.insert(vs[i].first);
      ans = max<long long>(tot + 1ll * ts.size() * ts.size(), ans);
    }
  }
  cout << ans << endl;
}

int main() {
  long long N;
  scanf("%lld", &N);
  long long K;
  scanf("%lld", &K);
  std::vector<long long> t(N);
  std::vector<long long> d(N);
  for (int i = 0; i < N; i++) {
    scanf("%lld", &t[i]);
    scanf("%lld", &d[i]);
  }
  solve(N, K, std::move(t), std::move(d));
  return 0;
}


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

using namespace std;

void solve(long long N, long long K, std::vector<long long> t,
           std::vector<long long> d) {
  vector<pair<long long, long long>> vs;
  for (int i = 0; i < N; i++) {
    vs.push_back({d[i], t[i]});
  }
  sort(vs.rbegin(), vs.rend());
  for (int i = 0; i < vs.size(); i++)
    swap(vs[i].first, vs[i].second);
  long long tot = 0;
  set<long long> ts;
  vector<long long> remCand;
  for (int i = 0; i < K; i++) {
    if (!ts.count(vs[i].first)) {
      ts.insert(vs[i].first);
    } else {
      remCand.push_back(vs[i].second);
    }
    tot += vs[i].second;
  }
  sort(remCand.rbegin(), remCand.rend());

  long long ans = tot + 1ll * ts.size() * ts.size();
  for (int i = K; i < N; i++) {
    if (!ts.count(vs[i].first) && remCand.size() > 0) {
      tot += vs[i].second;
      tot -= remCand.back();
      remCand.pop_back();
      ts.insert(vs[i].first);
      ans = max<long long>(tot + 1ll * ts.size() * ts.size(), ans);
    }
  }
  cout << ans << endl;
}

int main() {
  long long N;
  scanf("%lld", &N);
  long long K;
  scanf("%lld", &K);
  std::vector<long long> t(N);
  std::vector<long long> d(N);
  for (int i = 0; i < N; i++) {
    scanf("%lld", &t[i]);
    scanf("%lld", &d[i]);
  }
  solve(N, K, std::move(t), std::move(d));
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a) for (int i = 0; i < (a); i++)
#define ALL(a) (a).begin(), (a).end()
typedef long long ll;
typedef pair<int, int> P;
const int INF = 1e9;
const long long LINF = 1e18;
const long long MOD = 1e9 + 7;

signed main() {
  int n, k;
  cin >> n >> k;
  vector<pair<int, ll>> v(n);
  REP(i, n) { cin >> v[i].second >> v[i].first; }
  sort(ALL(v));
  reverse(ALL(v));
  ll now = 0;
  set<int> s;
  multiset<ll> used;
  for (int i = 0; i < k; i++) {
    now += v[i].first;
    if (s.count(v[i].second))
      used.insert(v[i].first);
    else
      s.insert(v[i].second);
  }
  now += s.size() * s.size();
  ll ans = now;
  for (int i = k; i < n; i++) {
    if (!s.count(v[i].second)) {
      now -= *used.begin();
      used.erase(used.begin());
      s.insert(v[i].second);
      now += v[i].first;
      now += s.size() * s.size() - (s.size() - 1) * (s.size() - 1);
      ans = max(ans, now);
    }
  }
  cout << ans << endl;
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a) for (int i = 0; i < (a); i++)
#define ALL(a) (a).begin(), (a).end()
typedef long long ll;
typedef pair<int, int> P;
const int INF = 1e9;
const long long LINF = 1e18;
const long long MOD = 1e9 + 7;

signed main() {
  int n, k;
  cin >> n >> k;
  vector<pair<int, ll>> v(n);
  REP(i, n) { cin >> v[i].second >> v[i].first; }
  sort(ALL(v));
  reverse(ALL(v));
  ll now = 0;
  set<int> s;
  multiset<ll> used;
  for (int i = 0; i < k; i++) {
    now += v[i].first;
    if (s.count(v[i].second))
      used.insert(v[i].first);
    else
      s.insert(v[i].second);
  }
  now += s.size() * s.size();
  ll ans = now;
  for (int i = k; i < n; i++) {
    if (used.size() == 0)
      break;
    if (!s.count(v[i].second)) {
      now -= *used.begin();
      used.erase(used.begin());
      s.insert(v[i].second);
      now += v[i].first;
      now += s.size() * s.size() - (s.size() - 1) * (s.size() - 1);
      ans = max(ans, now);
    }
  }
  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 pair<int, int> P;
#define rep(i, m, n) for (int i = m; i < n; i++)
ll mod = 1e9 + 7;

int main() {
  int n, k;
  cin >> n >> k;
  vector<int> v[n];
  rep(i, 0, n) {
    int a, d;
    cin >> a >> d;
    v[a - 1].push_back(d);
  }
  rep(i, 0, n) { sort(v[i].begin(), v[i].end(), greater<int>()); }
  priority_queue<int> t;
  priority_queue<int> f;
  rep(i, 0, n) {
    int l = v[i].size();
    rep(j, 0, l) {
      if (j == 0)
        t.push(v[i][j]);
      else
        f.push(v[i][j]);
    }
  }
  priority_queue<int> tt;
  priority_queue<int, vector<int>, greater<int>> ff;
  ll sum = 0;
  ll size = 0;
  rep(i, 0, k) {
    if (t.empty()) {
      ff.push(f.top());
      sum += f.top();
      f.pop();
    } else if (t.top() >= f.top()) {
      tt.push(t.top());
      sum += t.top();
      size++;
      t.pop();
    } else {
      ff.push(f.top());
      sum += f.top();
      f.pop();
    }
  }
  ll ans = sum + size * size;
  while (t.size() && ff.size()) {
    tt.push(t.top());
    sum += t.top();
    size++;
    t.pop();
    sum -= ff.top();
    ff.pop();
    ans = max(ans, sum + size * size);
  }
  cout << ans << endl;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> P;
#define rep(i, m, n) for (int i = m; i < n; i++)
ll mod = 1e9 + 7;

int main() {
  int n, k;
  cin >> n >> k;
  vector<int> v[n];
  rep(i, 0, n) {
    int a, d;
    cin >> a >> d;
    v[a - 1].push_back(d);
  }
  rep(i, 0, n) { sort(v[i].begin(), v[i].end(), greater<int>()); }
  priority_queue<int> t;
  priority_queue<int> f;
  rep(i, 0, n) {
    int l = v[i].size();
    rep(j, 0, l) {
      if (j == 0)
        t.push(v[i][j]);
      else
        f.push(v[i][j]);
    }
  }
  priority_queue<int> tt;
  priority_queue<int, vector<int>, greater<int>> ff;
  ll sum = 0;
  ll size = 0;
  rep(i, 0, k) {
    if (t.empty()) {
      ff.push(f.top());
      sum += f.top();
      f.pop();
    } else if (f.empty()) {
      tt.push(t.top());
      sum += t.top();
      size++;
      t.pop();
    } else if (t.top() >= f.top()) {
      tt.push(t.top());
      sum += t.top();
      size++;
      t.pop();
    } else {
      ff.push(f.top());
      sum += f.top();
      f.pop();
    }
  }
  ll ans = sum + size * size;
  while (t.size() && ff.size()) {
    tt.push(t.top());
    sum += t.top();
    size++;
    t.pop();
    sum -= ff.top();
    ff.pop();
    ans = max(ans, sum + size * size);
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<ll, ll>;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define ve vector

int main() {
  ll n, k;
  cin >> n >> k;
  ve<P> dt(n);
  ve<ll> fl(n, 0);
  rep(i, n) {
    cin >> dt[i].second >> dt[i].first;
    dt[i].second--;
  }
  sort(dt.rbegin(), dt.rend());
  ll ans = 0;
  ll cnt = 0;
  ve<P> kouho;
  rep(i, k) {
    ans += dt[i].first;
    if (fl[dt[i].second] == 0)
      cnt++;
    else
      kouho.push_back(dt[i]);
    fl[dt[i].second] += 1;
  }
  ans += cnt * cnt;
  ve<ll> x;
  x.push_back(ans);
  ll i = k;
  while (i < n) {
    if (fl[dt[i].second] != 0) {
      i++;
      continue;
    }
    sort(kouho.begin(), kouho.end());
    ans -= cnt * cnt;
    ans -= kouho[0].first;
    kouho.erase(kouho.begin());
    cnt++;
    ans += cnt * cnt;
    ans += dt[i].first;
    fl[dt[i].second] = 1;
    x.push_back(ans);
    i++;
  }
  sort(x.rbegin(), x.rend());
  cout << x[0] << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<ll, ll>;
#define rep(i, n) for (ll i = 0; i < (ll)(n); i++)
#define ve vector

int main() {
  ll n, k;
  cin >> n >> k;
  ve<P> dt(n);
  ve<ll> fl(n, 0);
  rep(i, n) {
    cin >> dt[i].second >> dt[i].first;
    dt[i].second--;
  }
  sort(dt.rbegin(), dt.rend());
  ll ans = 0;
  ll cnt = 0;
  ve<P> kouho;
  rep(i, k) {
    ans += dt[i].first;
    if (fl[dt[i].second] == 0)
      cnt++;
    else
      kouho.push_back(dt[i]);
    fl[dt[i].second] += 1;
  }
  ans += cnt * cnt;
  ve<ll> x;
  x.push_back(ans);
  ll i = k;
  while (i < n && !kouho.empty()) {
    if (fl[dt[i].second] != 0) {
      i++;
      continue;
    }
    sort(kouho.begin(), kouho.end());
    ans -= cnt * cnt;
    ans -= kouho[0].first;
    kouho.erase(kouho.begin());
    cnt++;
    ans += cnt * cnt;
    ans += dt[i].first;
    fl[dt[i].second] = 1;
    x.push_back(ans);
    i++;
  }
  sort(x.rbegin(), x.rend());
  cout << x[0] << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<pair<int, int>> sushi(N);
  for (int i = 0; i < N; ++i) {
    int t, d;
    cin >> t >> d;
    --t;
    sushi[i].first = d;
    sushi[i].second = t;
  }

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

  long long deliciousness = 0, variety = 0;
  vector<bool> is_used(N, false);
  priority_queue<int, vector<int>, greater<int>> pq;
  for (int i = 0; i < K; ++i) {
    deliciousness += sushi[i].first;
    if (!is_used[sushi[i].second]) {
      is_used[sushi[i].second] = true;
      ++variety;
    } else
      pq.push(sushi[i].first);
  }

  long long ans = deliciousness + variety * variety;
  for (int i = K; i < N; ++i) {
    if (!is_used[sushi[i].second]) {
      int tmp = pq.top();
      pq.pop();
      is_used[sushi[i].second] = true;
      deliciousness += sushi[i].first - tmp;
      ++variety;
    }
    ans = max(ans, deliciousness + variety * variety);
  }
  cout << ans << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<pair<int, int>> sushi(N);
  for (int i = 0; i < N; ++i) {
    int t, d;
    cin >> t >> d;
    --t;
    sushi[i].first = d;
    sushi[i].second = t;
  }

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

  long long deliciousness = 0, variety = 0;
  vector<bool> is_used(N, false);
  priority_queue<int, vector<int>, greater<int>> pq;
  for (int i = 0; i < K; ++i) {
    deliciousness += sushi[i].first;
    if (!is_used[sushi[i].second]) {
      is_used[sushi[i].second] = true;
      ++variety;
    } else
      pq.push(sushi[i].first);
  }

  long long ans = deliciousness + variety * variety;
  for (int i = K; i < N; ++i) {
    if (!is_used[sushi[i].second] and !pq.empty()) {
      int tmp = pq.top();
      pq.pop();
      is_used[sushi[i].second] = true;
      deliciousness += sushi[i].first - tmp;
      ++variety;
    }
    ans = max(ans, deliciousness + variety * variety);
  }
  cout << ans << endl;

  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define pb push_back
#define ll long long
#define mp make_pair
#define inf 1000000007
using namespace std;
int used[114514];
signed main() {
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> sus;
  for (int i = 0; i < n; i++) {
    int x, y;
    cin >> x >> y;
    sus.pb(mp(y, x));
  }
  sort(sus.begin(), sus.end());
  reverse(sus.begin(), sus.end());
  vector<ll> fir;
  priority_queue<ll> dup;
  ll sum = 0;
  for (int i = 0; i < k; i++) {
    if (used[sus[i].second] == 0) {
      used[sus[i].second] = 1;
      fir.push_back(sus[i].first);
      sum += sus[i].first;
    } else {
      dup.push(-sus[i].first);
      sum += sus[i].first;
    }
  }
  vector<ll> anss;
  anss.pb(sum + fir.size() * fir.size());
  for (int i = k; i < sus.size(); i++) {
    if (used[sus[i].second] == 0) {
      used[sus[i].second] = 1;
      sum -= -dup.top();
      dup.pop();
      fir.push_back(-sus[i].first);
      sum += sus[i].first;
      anss.pb(sum + fir.size() * fir.size());
    }
  }
  sort(anss.begin(), anss.end());
  cout << anss[anss.size() - 1] << endl;
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
#define pb push_back
#define ll long long
#define mp make_pair
#define inf 1000000007
using namespace std;
int used[114514];
signed main() {
  int n, k;
  cin >> n >> k;
  vector<pair<int, int>> sus;
  for (int i = 0; i < n; i++) {
    int x, y;
    cin >> x >> y;
    sus.pb(mp(y, x));
  }
  sort(sus.begin(), sus.end());
  reverse(sus.begin(), sus.end());
  vector<ll> fir;
  priority_queue<ll> dup;
  ll sum = 0;
  for (int i = 0; i < k; i++) {
    if (used[sus[i].second] == 0) {
      used[sus[i].second] = 1;
      fir.push_back(sus[i].first);
      sum += sus[i].first;
    } else {
      dup.push(-sus[i].first);
      sum += sus[i].first;
    }
  }
  vector<ll> anss;
  anss.pb(sum + fir.size() * fir.size());
  for (int i = k; i < sus.size(); i++) {
    if (used[sus[i].second] == 0 && !dup.empty()) {
      used[sus[i].second] = 1;
      sum -= -dup.top();
      dup.pop();
      fir.push_back(-sus[i].first);
      sum += sus[i].first;
      anss.pb(sum + fir.size() * fir.size());
    }
  }
  sort(anss.begin(), anss.end());
  cout << anss[anss.size() - 1] << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

using namespace std;

const int N = 1e5 + 5;

typedef long long ll;

int u, p, d, n, tc, q, k, cnt;

char s[N];

ll ans;

vector<int> topping[N];

pair<int, int> a[N];

multiset<int> m, kind;

vector<pair<int, int>> v;

int main() {
  //	freopen("test.txt","r",stdin);
  scanf("%d %d", &n, &k);
  for (int i = 0; i < n; i++)
    scanf("%d%d", &a[i].first, &a[i].second),
        topping[--a[i].first].push_back(a[i].second);
  ll del = 0;
  for (int i = 0; i < n; i++) {
    if (topping[i].size()) {
      sort(topping[i].begin(), topping[i].end(), greater<int>());
      v.push_back({topping[i][0], i});
    }
  }
  sort(v.begin(), v.end(), greater<pair<int, int>>());
  for (int ii = 0; ii < v.size() && ii < k; ii++) {
    int i = v[ii].second;
    del += topping[i][0];
    if (m.size() + ii + 1 > k) {
      del -= *m.begin();
      m.erase(m.begin());
    }
    for (int j = 1; j < topping[i].size(); j++) {
      if ((ii + 1 + (int)m.size()) == k) {
        if (topping[i][j] > *m.begin()) {
          del -= *m.begin();
          m.erase(m.begin());
        } else
          break;
      }
      del += topping[i][j];
      m.insert(topping[i][j]);
    }
    ll sz = ii + 1;
    ans = max(ans, sz * sz + del);
  }
  printf("%lld\n", ans);
}


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

using namespace std;

const int N = 1e5 + 5;

typedef long long ll;

int u, p, d, n, tc, q, k, cnt;

char s[N];

ll ans;

vector<int> topping[N];

pair<int, int> a[N];

multiset<int> m, kind;

vector<pair<int, int>> v;

int main() {
  //	freopen("test.txt","r",stdin);
  scanf("%d %d", &n, &k);
  for (int i = 0; i < n; i++)
    scanf("%d%d", &a[i].first, &a[i].second),
        topping[--a[i].first].push_back(a[i].second);
  ll del = 0;
  for (int i = 0; i < n; i++) {
    if (topping[i].size()) {
      sort(topping[i].begin(), topping[i].end(), greater<int>());
      v.push_back({topping[i][0], i});
    }
  }
  sort(v.begin(), v.end(), greater<pair<int, int>>());
  for (int ii = 0; ii < v.size() && ii < k; ii++) {
    int i = v[ii].second;
    del += topping[i][0];
    if (m.size() + ii + 1 > k) {
      del -= *m.begin();
      m.erase(m.begin());
    }
    for (int j = 1; j < topping[i].size(); j++) {
      if ((ii + 1 + (int)m.size()) == k) {
        if (!m.size())
          break;
        if (topping[i][j] > *m.begin()) {
          del -= *m.begin();
          m.erase(m.begin());
        } else
          break;
      }
      del += topping[i][j];
      m.insert(topping[i][j]);
    }
    ll sz = ii + 1;
    ans = max(ans, sz * sz + del);
  }
  printf("%lld\n", ans);
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>
#define ll long long
using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<vector<ll>> v(N + 1);
  ll sum = 0;
  for (int i = 0; i < N; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    v[t].push_back(d);
    sum += d;
  }
  priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<pair<ll, int>>>
      pq_g;
  priority_queue<ll, vector<ll>, greater<ll>> pq_f;
  vector<int> cnt(N + 1, 0);
  ll num = 0;
  for (int i = 1; i <= N; ++i) {
    if (v[i].size() > 0) {
      ++num;
      sort(v[i].begin(), v[i].end(), greater<ll>());
      if (v[i].size() == 1)
        pq_f.push(v[i][0]);
      else {
        for (int j = 1; j < v[i].size(); ++j)
          pq_g.push(make_pair(v[i][j], i));
      }
    }
  }
  for (int i = N; i > K; --i) {
    if (pq_f.size() == 0) {
      sum -= pq_g.top().first;
      ++cnt[pq_g.top().second];
      if (cnt[pq_g.top().second] + 1 == v[pq_g.top().second].size()) {
        pq_f.push(v[pq_g.top().second][0]);
      }
      pq_g.pop();
    } else if (pq_g.top().first < 2 * num - 1 + pq_f.top()) {
      sum -= pq_g.top().first;
      ++cnt[pq_g.top().second];
      if (cnt[pq_g.top().second] + 1 == v[pq_g.top().second].size()) {
        pq_f.push(v[pq_g.top().second][0]);
      }
      pq_g.pop();
    } else {
      sum -= pq_f.top();
      --num;
      pq_f.pop();
    }
  }
  cout << sum + num * num << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>
#define ll long long
using namespace std;

int main() {
  int N, K;
  cin >> N >> K;
  vector<vector<ll>> v(N + 1);
  ll sum = 0;
  for (int i = 0; i < N; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    v[t].push_back(d);
    sum += d;
  }
  priority_queue<pair<ll, int>, vector<pair<ll, int>>, greater<pair<ll, int>>>
      pq_g;
  priority_queue<ll, vector<ll>, greater<ll>> pq_f;
  vector<int> cnt(N + 1, 0);
  ll num = 0;
  for (int i = 1; i <= N; ++i) {
    if (v[i].size() > 0) {
      ++num;
      sort(v[i].begin(), v[i].end(), greater<ll>());
      if (v[i].size() == 1)
        pq_f.push(v[i][0]);
      else {
        for (int j = 1; j < v[i].size(); ++j)
          pq_g.push(make_pair(v[i][j], i));
      }
    }
  }
  for (int i = N; i > K; --i) {
    if (pq_f.size() == 0) {
      sum -= pq_g.top().first;
      ++cnt[pq_g.top().second];
      if (cnt[pq_g.top().second] + 1 == v[pq_g.top().second].size()) {
        pq_f.push(v[pq_g.top().second][0]);
      }
      pq_g.pop();
    } else if (pq_g.size() == 0) {
      sum -= pq_f.top();
      --num;
      pq_f.pop();
    } else if (pq_g.top().first < 2 * num - 1 + pq_f.top()) {
      sum -= pq_g.top().first;
      ++cnt[pq_g.top().second];
      if (cnt[pq_g.top().second] + 1 == v[pq_g.top().second].size()) {
        pq_f.push(v[pq_g.top().second][0]);
      }
      pq_g.pop();
    } else {
      sum -= pq_f.top();
      --num;
      pq_f.pop();
    }
  }
  cout << sum + num * num << endl;
  return 0;
}
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
n, k = map(int, input().split())
td = list(
    sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1])
)
td2 = []
eaten = set()
dup = 0
score = k * k
i = 0
j = 0
while len(eaten) + dup < k:
    if td[i][0] in eaten:
        td2.append(td[i])
        i += 1
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if i < n and (j >= len(td2) or td[i][1] >= td2[j][1] - diff):
            t, d = td[i]
            score += d
            eaten.add(t)
            i += 1
        else:
            t, d = td2[j]
            score += d - diff
            dup += 1
            j += 1
print(score)


Fixed submission:
n, k = map(int, input().split())
td = list(
    sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1])
)
td2 = []
eaten = set()
dup = 0
score = k * k
i = 0
j = 0
while len(eaten) + dup < k:
    if i < n and td[i][0] in eaten:
        td2.append(td[i])
        i += 1
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if i < n and (j >= len(td2) or td[i][1] >= td2[j][1] - diff):
            t, d = td[i]
            score += d
            eaten.add(t)
            i += 1
        else:
            t, d = td2[j]
            score += d - diff
            dup += 1
            j += 1
print(score)
Problem p03148 (C++)

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

typedef long long ll;
const int M = 10010;

struct node {
  ll v, x;
} a[M];

ll vis[M];

bool cmp(node x, node y) { return x.v > y.v; }

int n, k;
ll cnt, tp, res, ans;
int main() {
  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    int a1, b;
    cin >> a1 >> b;
    a[i].x = a1;
    a[i].v = b;
  }
  sort(a + 1, a + n + 1, cmp);
  stack<ll> s;
  for (int i = 1; i <= n; i++) {
    if (cnt < k) {
      if (vis[a[i].x] == 0) {
        vis[a[i].x] = 1;
        tp++;
      } else {
        s.push(a[i].v);
      }
      res += a[i].v;
      cnt++;
      ans = max(ans, res + 1ll * tp * tp);
    } else {
      if (s.empty())
        break;
      if (vis[a[i].x])
        continue;
      vis[a[i].x] = 1;
      tp++;
      res -= s.top();
      res += a[i].v;
      s.pop();
      ans = max(ans, res + tp * tp);
    }
  }
  printf("%lld\n", ans);
  return 0;
}

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

typedef long long ll;
const int M = 100010;

struct node {
  ll v, x;
} a[M];

ll vis[M];

bool cmp(node x, node y) { return x.v > y.v; }

int n, k;
ll cnt, tp, res, ans;
int main() {
  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    int a1, b;
    cin >> a1 >> b;
    a[i].x = a1;
    a[i].v = b;
  }
  sort(a + 1, a + n + 1, cmp);
  stack<ll> s;
  for (int i = 1; i <= n; i++) {
    if (cnt < k) {
      if (vis[a[i].x] == 0) {
        vis[a[i].x] = 1;
        tp++;
      } else {
        s.push(a[i].v);
      }
      res += a[i].v;
      cnt++;
      ans = max(ans, res + 1ll * tp * tp);
    } else {
      if (s.empty())
        break;
      if (vis[a[i].x])
        continue;
      vis[a[i].x] = 1;
      tp++;
      res -= s.top();
      res += a[i].v;
      s.pop();
      ans = max(ans, res + tp * tp);
    }
  }
  printf("%lld\n", ans);
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#include <cinttypes>

using namespace std;
using ll = long long;
const int NMAX = 100010;

int used[NMAX];
pair<int, int> dt[NMAX];
ll f[NMAX];

int main() {
  ios::sync_with_stdio(false);
  int N, K;
  scanf("%d %d", &N, &K);
  // vector<pair<int, int>> dt(N);
  // vector<int> used(N+1, 0);
  fill(used, used + N + 1, 0);
  // vector<ll> f(N+1, -1LL);
  fill(f, f + N + 1, -1);
  for (int i = 0; i < N; i++)
    scanf("%d %d", &(dt[i].second), &(dt[i].first)); // first = d, second = t;

  sort(dt, dt + N, greater<pair<int, int>>());

  ll goodness = 0;
  priority_queue<pair<int, int>, vector<pair<int, int>>,
                 greater<pair<int, int>>>
      que;
  int used_cnt = 0;
  for (int i = 0; i < K; i++) {
    que.push(dt[i]);
    goodness += dt[i].first;
    if (used[dt[i].second] == false) {
      used_cnt++;
    }
    used[dt[i].second]++;
  }

  f[used_cnt] = goodness;

  for (int i = K; i < N; i++) {
    if (used[dt[i].second] > 0)
      continue;
    while (used[que.top().second] == 1) {
      que.pop();
    }
    if (que.empty()) {
      break;
    }
    pair<int, int> q = que.top();
    que.pop();
    goodness -= q.first;
    goodness += dt[i].first;
    used[dt[i].second]++;
    used_cnt++;
    f[used_cnt] = goodness;
  }

  ll ans = 0;
  for (int i = 1; i <= N; i++) {
    if (f[i] == -1)
      continue;
    ans = max(ans, f[i] + (ll)i * (ll)i);
  }
  cout << ans << '\n';
  return 0;
}

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

using namespace std;
using ll = long long;
const int NMAX = 100010;

int used[NMAX];
pair<int, int> dt[NMAX];
ll f[NMAX];

int main() {
  ios::sync_with_stdio(false);
  int N, K;
  scanf("%d %d", &N, &K);
  // vector<pair<int, int>> dt(N);
  // vector<int> used(N+1, 0);
  fill(used, used + N + 1, 0);
  // vector<ll> f(N+1, -1LL);
  fill(f, f + N + 1, -1);
  for (int i = 0; i < N; i++)
    scanf("%d %d", &(dt[i].second), &(dt[i].first)); // first = d, second = t;

  sort(dt, dt + N, greater<pair<int, int>>());

  ll goodness = 0;
  priority_queue<pair<int, int>, vector<pair<int, int>>,
                 greater<pair<int, int>>>
      que;
  int used_cnt = 0;
  for (int i = 0; i < K; i++) {
    que.push(dt[i]);
    goodness += dt[i].first;
    if (used[dt[i].second] == false) {
      used_cnt++;
    }
    used[dt[i].second]++;
  }

  f[used_cnt] = goodness;

  for (int i = K; i < N; i++) {
    if (used[dt[i].second] > 0)
      continue;
    while (!que.empty() && used[que.top().second] == 1) {
      que.pop();
    }
    if (que.empty()) {
      break;
    }
    pair<int, int> q = que.top();
    que.pop();
    goodness -= q.first;
    goodness += dt[i].first;
    used[dt[i].second]++;
    used_cnt++;
    f[used_cnt] = goodness;
  }

  ll ans = 0;
  for (int i = 1; i <= N; i++) {
    if (f[i] == -1)
      continue;
    ans = max(ans, f[i] + (ll)i * (ll)i);
  }
  cout << ans << '\n';
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;
typedef long long ll;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef pair<ll, ll> P;
typedef vector<P> vpl;
#define rep(i, n) for (ll i = 0; i < (n); i++)
#define REP(i, a, b) for (ll i = (a); i < (b); i++)
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
const int inf = 1 << 30;
const ll linf = 1LL << 62;
const int mod = 1e9 + 7;
const int MAX = 510000;
const int V = 100005;

int main() {
  ll n, k;
  cin >> n >> k;
  priority_queue<P> pq;
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    t--;
    pq.emplace(P(d, t));
  }
  vl sushi(n, 0);
  ll sum = 0;
  ll x = 0;
  multiset<ll> mse;
  rep(i, k) {
    auto p = pq.top();
    pq.pop();
    sum += p.first;
    ll u = p.second;
    if (sushi[u] == 0) {
      sushi[u]++;
      x++;
    } else
      mse.insert(p.first);
  }
  ll ans = sum + x * x;
  while (x < k && !mse.empty()) {
    auto p = pq.top();
    pq.pop();
    ll u = p.second;
    if (sushi[u] > 0)
      continue;
    sushi[u]++;
    sum -= *mse.begin();
    mse.erase(mse.begin());
    sum += p.first;
    x++;
    ans = max(ans, sum + x * x);
  }
  cout << ans << endl;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;
typedef long long ll;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef pair<ll, ll> P;
typedef vector<P> vpl;
#define rep(i, n) for (ll i = 0; i < (n); i++)
#define REP(i, a, b) for (ll i = (a); i < (b); i++)
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
const int inf = 1 << 30;
const ll linf = 1LL << 62;
const int mod = 1e9 + 7;
const int MAX = 510000;
const int V = 100005;

int main() {
  ll n, k;
  cin >> n >> k;
  priority_queue<P> pq;
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    t--;
    pq.emplace(P(d, t));
  }
  vl sushi(n, 0);
  ll sum = 0;
  ll x = 0;
  multiset<ll> mse;
  rep(i, k) {
    auto p = pq.top();
    pq.pop();
    sum += p.first;
    ll u = p.second;
    if (sushi[u] == 0) {
      sushi[u]++;
      x++;
    } else
      mse.insert(p.first);
  }
  ll ans = sum + x * x;
  while (x < k && !mse.empty() && !pq.empty()) {
    auto p = pq.top();
    pq.pop();
    ll u = p.second;
    if (sushi[u] > 0)
      continue;
    sushi[u]++;
    sum -= *mse.begin();
    mse.erase(mse.begin());
    sum += p.first;
    x++;
    ans = max(ans, sum + x * x);
  }
  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 < int(n); ++i)
#define all(a) (a).begin(), (a).end()
typedef long long ll;
using P = pair<ll, ll>;

template <typename T> vector<T> cumulative_sum(vector<T> &v) {
  vector<T> sum(v.size() + 1);
  rep(i, v.size()) { sum[i + 1] = sum[i] + v[i]; }
  return sum;
}

int main(void) {
  int n;
  ll K;
  cin >> n >> K;
  vector<P> vp(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    vp[i] = {d, t};
  }
  sort(all(vp), greater<P>());
  map<ll, bool> used;
  vector<ll> top, rem;
  rep(i, n) {
    ll t = vp[i].second;
    ll d = vp[i].first;
    if (used[t]) {
      rem.push_back(d);
    } else {
      used[t] = true;
      top.push_back(d);
    }
  }
  auto top_sum = cumulative_sum(top);
  auto rem_sum = cumulative_sum(rem);
  ll res = 0;
  for (ll x = 1; x <= K; x++) {
    if (K < x)
      continue;
    if (K - x >= (ll)rem_sum.size())
      continue;
    res = max(res, top_sum[x] + rem_sum[K - x] + x * x);
  }
  cout << res << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < int(n); ++i)
#define all(a) (a).begin(), (a).end()
typedef long long ll;
using P = pair<ll, ll>;

template <typename T> vector<T> cumulative_sum(vector<T> &v) {
  vector<T> sum(v.size() + 1);
  rep(i, v.size()) { sum[i + 1] = sum[i] + v[i]; }
  return sum;
}

int main(void) {
  int n;
  ll K;
  cin >> n >> K;
  vector<P> vp(n);
  rep(i, n) {
    ll t, d;
    cin >> t >> d;
    vp[i] = {d, t};
  }
  sort(all(vp), greater<P>());
  map<ll, bool> used;
  vector<ll> top, rem;
  rep(i, n) {
    ll t = vp[i].second;
    ll d = vp[i].first;
    if (used[t]) {
      rem.push_back(d);
    } else {
      used[t] = true;
      top.push_back(d);
    }
  }
  auto top_sum = cumulative_sum(top);
  auto rem_sum = cumulative_sum(rem);
  ll res = 0;
  for (ll x = 1; x <= K; x++) {
    if (x >= (ll)top_sum.size())
      continue;
    if (K < x)
      continue;
    if (K - x >= (ll)rem_sum.size())
      continue;
    res = max(res, top_sum[x] + rem_sum[K - x] + x * x);
  }
  cout << res << endl;
  return 0;
}
Problem p03148 (C++)

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

typedef long long ll;
typedef unsigned long long ull;

int main() {
  ll N, K;
  cin >> N >> K;
  vector<array<ll, 2>> td(N);
  for (int i = 0; i < N; i++) {
    cin >> td[i][1] >> td[i][0];
  }
  sort(td.begin(), td.end(), greater<array<ll, 2>>());

  ll res = 0;
  ll vsum = 0;
  set<ll> kind;
  stack<ll> unused;
  for (int i = 0; i < K; i++) {
    auto t = td[i][1];
    auto d = td[i][0];
    vsum += d;
    ;
    if (kind.count(t) > 0)
      unused.push(d);
    else
      kind.insert(t);
  }
  res = vsum + kind.size() * kind.size();
  for (int i = K; i < N; i++) {
    auto t = td[i][1];
    auto d = td[i][0];
    if (kind.count(t) > 0)
      continue;
    kind.insert(t);
    vsum += d;
    vsum -= unused.top();
    unused.pop();
    res = max(res, (ll)(vsum + kind.size() * kind.size()));
  }
  cout << res << endl;

  return 0;
}

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

typedef long long ll;
typedef unsigned long long ull;

int main() {
  ll N, K;
  cin >> N >> K;
  vector<array<ll, 2>> td(N);
  for (int i = 0; i < N; i++) {
    cin >> td[i][1] >> td[i][0];
  }
  sort(td.begin(), td.end(), greater<array<ll, 2>>());

  ll res = 0;
  ll vsum = 0;
  set<ll> kind;
  stack<ll> unused;
  for (int i = 0; i < K; i++) {
    auto t = td[i][1];
    auto d = td[i][0];
    vsum += d;
    ;
    if (kind.count(t) > 0)
      unused.push(d);
    else
      kind.insert(t);
  }
  res = vsum + kind.size() * kind.size();
  for (int i = K; i < N; i++) {
    auto t = td[i][1];
    auto d = td[i][0];
    if (kind.count(t) > 0)
      continue;
    kind.insert(t);
    vsum += d;
    if (unused.size() == 0)
      break;
    vsum -= unused.top();
    unused.pop();
    res = max(res, (ll)(vsum + kind.size() * kind.size()));
  }
  cout << res << endl;

  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string.h>
#include <string>
#include <utility>
#include <vector>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using piii = pair<pii, pii>;
const int INF = 1e9 + 7;
int main() {
  int N, K;
  cin >> N >> K;
  priority_queue<ll> kind[N];
  for (int i = 0; i < N; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    kind[t - 1].push(d);
  }
  priority_queue<ll> best, remain;
  // best:各種類で一番おいしいやつを集めたもの
  // remain:bestの余り
  for (int i = 0; i < N; ++i) {
    if (!kind[i].empty()) {
      best.push(kind[i].top());
      kind[i].pop();
      while (!kind[i].empty())
        ;
      {
        remain.push(kind[i].top());
        kind[i].pop();
      }
    }
  }
  vector<ll> bsum = {0};
  // bestの累積和(bsum[i]=bestの先頭からi個の総和)
  while (!best.empty()) {
    bsum.push_back(bsum.back() + best.top());
    best.pop();
  }
  vector<ll> rsum = {0};
  // remainの累積和
  while (!remain.empty()) {
    rsum.push_back(rsum.back() + remain.top());
    remain.pop();
  }
  ll ans = 0;
  for (ll x = 1; x <= K; ++x) {
    // 種類が多すぎるケース
    if (x >= bsum.size())
      break;
    // 種類が少なすぎて余り物が足りないケース
    if (K - x >= rsum.size())
      continue;
    // bestからx個、remainからK-x個貪欲に選ぶ
    ans = max(ans, bsum[x] + rsum[K - x] + x * x);
  }
  cout << ans << endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <deque>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string.h>
#include <string>
#include <utility>
#include <vector>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using piii = pair<pii, pii>;
const int INF = 1e9 + 7;
int main() {
  int N, K;
  cin >> N >> K;
  priority_queue<ll> kind[N];
  for (int i = 0; i < N; ++i) {
    int t;
    ll d;
    cin >> t >> d;
    kind[t - 1].push(d);
  }
  priority_queue<ll> best, remain;
  // best:各種類で一番おいしいやつを集めたもの
  // remain:bestの余り
  for (int i = 0; i < N; ++i) {
    if (!kind[i].empty()) {
      best.push(kind[i].top());
      kind[i].pop();
      while (!kind[i].empty()) {
        remain.push(kind[i].top());
        kind[i].pop();
      }
    }
  }
  vector<ll> bsum = {0};
  // bestの累積和(bsum[i]=bestの先頭からi個の総和)
  while (!best.empty()) {
    bsum.push_back(bsum.back() + best.top());
    best.pop();
  }
  vector<ll> rsum = {0};
  // remainの累積和
  while (!remain.empty()) {
    rsum.push_back(rsum.back() + remain.top());
    remain.pop();
  }
  ll ans = 0;
  for (ll x = 1; x <= K; ++x) {
    // 種類が多すぎるケース
    if (x >= bsum.size())
      break;
    // 種類が少なすぎて余り物が足りないケース
    if (K - x >= rsum.size())
      continue;
    // bestからx個、remainからK-x個貪欲に選ぶ
    ans = max(ans, bsum[x] + rsum[K - x] + x * x);
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Time Limit Exceeded):
#include "bits/stdc++.h"
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
typedef long long int ll;
typedef pair<ll, ll> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<P> v;
  rep(i, n) {
    int t;
    ll d;
    cin >> t >> d;
    v.emplace_back(d, t - 1);
  }
  sort(v.begin(), v.end(), greater<>());
  ll sum = 0, type = 0;
  vector<int> cnt(n, 0);
  rep(i, k) {
    sum += v[i].first;
    if (cnt[v[i].second]++ == 0)
      ++type;
  }
  ll ans = sum + type * type;
  for (int i = k, j = k - 1; i < n; ++i) {
    if (cnt[v[i].second] > 0)
      continue;
    ++cnt[v[i].second];
    sum += v[i].first;
    ++type;
    while (j >= 0) {
      if (cnt[v[j].second] > 1) {
        --cnt[v[j].second];
        sum -= v[j].first;
        chmax(ans, sum + type * type);
        --j;
        break;
      }
    }
    if (j < 0)
      break;
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include "bits/stdc++.h"
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
typedef long long int ll;
typedef pair<ll, ll> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<P> v;
  rep(i, n) {
    int t;
    ll d;
    cin >> t >> d;
    v.emplace_back(d, t - 1);
  }
  sort(v.begin(), v.end(), greater<>());
  ll sum = 0, type = 0;
  vector<int> cnt(n, 0);
  rep(i, k) {
    sum += v[i].first;
    if (cnt[v[i].second]++ == 0)
      ++type;
  }
  ll ans = sum + type * type;
  for (int i = k, j = k - 1; i < n; ++i) {
    if (cnt[v[i].second] > 0)
      continue;
    ++cnt[v[i].second];
    sum += v[i].first;
    ++type;
    while (j >= 0) {
      if (cnt[v[j].second] > 1) {
        --cnt[v[j].second];
        sum -= v[j].first;
        chmax(ans, sum + type * type);
        --j;
        break;
      }
      --j;
    }
    if (j < 0)
      break;
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include "bits/stdc++.h"
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
typedef long long int ll;
typedef pair<ll, int> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<P> sushi;
  rep(i, n) {
    int t;
    ll d;
    cin >> t >> d;
    sushi.emplace_back(d, t);
  }
  sort(sushi.rbegin(), sushi.rend());
  ll sum = 0, ans = 0, kind = 0;
  priority_queue<P> pq;
  vector<int> cnt(n + 1, 0);
  rep(i, k) {
    ++kind;
    if (cnt[sushi[i].second]++ > 0) {
      pq.push(P(-sushi[i].first, sushi[i].second));
      --kind;
    }
    sum += sushi[i].first;
  }
  ans = sum + kind * kind;
  for (int i = k; i < n; ++i) {
    if (cnt[sushi[i].second] > 0)
      continue;
    ll d = -pq.top().first;
    int t = pq.top().second;
    pq.pop();
    sum = sum - d + sushi[i].first;
    ++cnt[sushi[i].second];
    --cnt[t];
    ++kind;
    chmax(ans, sum + kind * kind);
  }
  cout << ans << endl;
  return 0;
}


Fixed submission:
#include "bits/stdc++.h"
#define rep(i, n) for (int i = 0; i < (n); ++i)
using namespace std;
typedef long long int ll;
typedef pair<ll, int> P;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int n, k;
  cin >> n >> k;
  vector<P> sushi;
  rep(i, n) {
    int t;
    ll d;
    cin >> t >> d;
    sushi.emplace_back(d, t);
  }
  sort(sushi.rbegin(), sushi.rend());
  ll sum = 0, ans = 0, kind = 0;
  priority_queue<P> pq;
  vector<int> cnt(n + 1, 0);
  rep(i, k) {
    ++kind;
    if (cnt[sushi[i].second]++ > 0) {
      pq.push(P(-sushi[i].first, sushi[i].second));
      --kind;
    }
    sum += sushi[i].first;
  }
  ans = sum + kind * kind;
  for (int i = k; i < n; ++i) {
    if (cnt[sushi[i].second] > 0)
      continue;
    if (pq.empty())
      break;
    ll d = -pq.top().first;
    int t = pq.top().second;
    pq.pop();
    sum = sum - d + sushi[i].first;
    ++cnt[sushi[i].second];
    --cnt[t];
    ++kind;
    chmax(ans, sum + kind * kind);
  }
  cout << ans << endl;
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <functional>
#include <stack>
#include <stdio.h>
#include <utility>
typedef long long ll;
using namespace std;

int main(void) {
  ll i, j, k, n, ans, now = 0, cnt = 0;
  scanf("%lld%lld", &n, &k);
  pair<ll, ll> s[n];
  bool f[n];
  for (i = 0; i < n; ++i)
    scanf("%lld%lld", &s[i].second, &s[i].first), s[i].second--, f[i] = true;
  sort(s, s + n, greater<pair<ll, ll>>());
  stack<ll> st;
  for (i = 0; i < k; ++i) {
    if (f[s[i].second]) {
      f[s[i].second] = false;
      cnt++;
    } else
      st.push(s[i].first);
    now += s[i].first;
  }
  ans = now + cnt * cnt;
  for (; i < n; ++i)
    if (f[s[i].second]) {
      now -= st.top();
      st.pop();
      now += s[i].first;
      cnt++;
      f[s[i].second] = false;
      if (now + cnt * cnt > ans)
        ans = now + cnt * cnt;
    }
  printf("%lld", ans);
  return 0;
}

Fixed submission:
#include <algorithm>
#include <functional>
#include <stack>
#include <stdio.h>
#include <utility>
typedef long long ll;
using namespace std;

int main(void) {
  ll i, j, k, n, ans, now = 0, cnt = 0;
  scanf("%lld%lld", &n, &k);
  pair<ll, ll> s[n];
  bool f[n];
  for (i = 0; i < n; ++i)
    scanf("%lld%lld", &s[i].second, &s[i].first), s[i].second--, f[i] = true;
  sort(s, s + n, greater<pair<ll, ll>>());
  stack<ll> st;
  for (i = 0; i < k; ++i) {
    if (f[s[i].second]) {
      f[s[i].second] = false;
      cnt++;
    } else
      st.push(s[i].first);
    now += s[i].first;
  }
  ans = now + cnt * cnt;
  for (; i < n && st.size(); ++i)
    if (f[s[i].second]) {
      now -= st.top();
      st.pop();
      now += s[i].first;
      cnt++;
      f[s[i].second] = false;
      if (now + cnt * cnt > ans)
        ans = now + cnt * cnt;
    }
  printf("%lld", ans);
  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define all(x) begin(x), end(x)
#define dbg(x) cerr << #x << " = " << x << endl
#define _ << ' ' <<
using namespace std;
using ll = long long;
using vi = vector<int>;

vector<ll> a[100001];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int k, n;
  cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    int t, d;
    cin >> t >> d;
    a[t].push_back(d);
  }
  vector<ll> bst, ot;
  for (int i = 1; i <= n; ++i) {
    if (a[i].empty())
      continue;
    sort(all(a[i]));
    reverse(all(a[i]));
    bst.push_back(a[i][0]);
    for (int j = 1; j < a[i].size(); ++j)
      ot.push_back(a[i][j]);
  }
  sort(all(bst));
  reverse(all(bst));
  sort(all(ot));
  reverse(all(ot));
  for (int i = 1; i < bst.size(); ++i)
    bst[i] += bst[i - 1];
  for (int i = 1; i < ot.size(); ++i)
    ot[i] += ot[i - 1];

  ll sol = 0;
  for (int i = 0; i < min((int)bst.size(), k); ++i) {
    ll att = bst[i] + 1ll * (i + 1) * (i + 1);
    if (k - i - 2 >= 0)
      att += ot[max(0, min((int)ot.size() - 1, k - i - 2))];
    sol = max(sol, att);
  }
  cout << sol;
}


Fixed submission:
#include <bits/stdc++.h>
#define all(x) begin(x), end(x)
#define dbg(x) cerr << #x << " = " << x << endl
#define _ << ' ' <<
using namespace std;
using ll = long long;
using vi = vector<int>;

vector<ll> a[100001];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int k, n;
  cin >> n >> k;
  for (int i = 0; i < n; ++i) {
    int t, d;
    cin >> t >> d;
    a[t].push_back(d);
  }
  vector<ll> bst, ot;
  for (int i = 1; i <= n; ++i) {
    if (a[i].empty())
      continue;
    sort(all(a[i]));
    reverse(all(a[i]));
    bst.push_back(a[i][0]);
    for (int j = 1; j < a[i].size(); ++j)
      ot.push_back(a[i][j]);
  }
  sort(all(bst));
  reverse(all(bst));
  sort(all(ot));
  reverse(all(ot));
  for (int i = 1; i < bst.size(); ++i)
    bst[i] += bst[i - 1];
  for (int i = 1; i < ot.size(); ++i)
    ot[i] += ot[i - 1];

  ll sol = 0;
  for (int i = 0; i < min((int)bst.size(), k); ++i) {
    ll att = bst[i] + 1ll * (i + 1) * (i + 1);
    if (k - i - 2 >= 0 && !ot.empty())
      att += ot[max(0, min((int)ot.size() - 1, k - i - 2))];
    sol = max(sol, att);
  }
  cout << sol;
}