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

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stdio.h>
#include <string>
#include <vector>
using namespace std;
#define int long long
int MOD = 1000000007;
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N, K;
  cin >> N >> K;
  vector<pair<int, int>> A(N);
  vector<int> B(N);
  int res = 0;
  for (int i = 0; i < N; i++) {
    cin >> A[i].second >> A[i].first;
  }
  sort(A.rbegin(), A.rend());
  int num = 0;
  int sum = 0;
  vector<int> cnt(N + 1, 0);
  vector<int> st;
  for (int i = 0; i < K; i++) {
    sum += A[i].first;
    // cerr << i << " " << A[i].second << endl;
    if (cnt[A[i].second] == 0) {
      cnt[A[i].second]++;
      num++;
    } else {
      cnt[A[i].second]++;
      st.push_back(i);
    }
  }
  // cerr << sum << " " << num << endl;
  res = max(res, sum + num * num);
  for (int i = K; i < N; i++) {
    if (cnt[A[i].second] == 0) {
      int j = st.back();
      st.pop_back();
      cnt[A[j].second]--;
      sum -= A[j].first;

      cnt[A[i].second]++;
      sum += A[i].first;
      num++;
      res = max(res, sum + num * num);
      // cerr << i << endl;
      // cerr << sum << " " << num << endl;
    }
  }

  cout << res << endl;
}

Fixed submission:
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stdio.h>
#include <string>
#include <vector>
using namespace std;
#define int long long
int MOD = 1000000007;
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N, K;
  cin >> N >> K;
  vector<pair<int, int>> A(N);
  vector<int> B(N);
  int res = 0;
  for (int i = 0; i < N; i++) {
    cin >> A[i].second >> A[i].first;
  }
  sort(A.rbegin(), A.rend());
  int num = 0;
  int sum = 0;
  vector<int> cnt(N + 1, 0);
  vector<int> st;
  for (int i = 0; i < K; i++) {
    sum += A[i].first;
    // cerr << i << " " << A[i].second << endl;
    if (cnt[A[i].second] == 0) {
      cnt[A[i].second]++;
      num++;
    } else {
      cnt[A[i].second]++;
      st.push_back(i);
    }
  }
  // cerr << sum << " " << num << endl;
  res = max(res, sum + num * num);
  for (int i = K; i < N; i++) {
    if (cnt[A[i].second] == 0) {
      if (st.size() == 0)
        break;
      int j = st.back();
      st.pop_back();
      cnt[A[j].second]--;
      sum -= A[j].first;

      cnt[A[i].second]++;
      sum += A[i].first;
      num++;
      res = max(res, sum + num * num);
      // cerr << i << endl;
      // cerr << sum << " " << num << endl;
    }
  }

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

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int MOD = 1e9 + 7;

signed main() {
  int N, K;
  cin >> N >> K;
  typedef pair<int, int> P;
  std::vector<P> v(N);
  for (int i = 0; i < N; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    v[i] = P(b, a);
  }
  sort(v.rbegin(), v.rend());
  int ans = 0;
  std::vector<int> usedc(N, 0);
  stack<int> sta;
  int res = 0;
  int color = 0;
  for (int i = 0; i < K; i++) {
    res += v[i].first;
    if (usedc[v[i].second]++ == 0)
      color++;
    else
      sta.push(v[i].first);
  }
  ans = res + color * color;
  for (int i = K; i < N; i++) {
    if (usedc[v[i].second]++)
      continue;
    int a = sta.top();
    sta.pop();
    res -= a;
    res += v[i].first;
    color++;
    ans = max(ans, res + color * color);
  }
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
#define int long long
using namespace std;
const int MOD = 1e9 + 7;

signed main() {
  int N, K;
  cin >> N >> K;
  typedef pair<int, int> P;
  std::vector<P> v(N);
  for (int i = 0; i < N; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    v[i] = P(b, a);
  }
  sort(v.rbegin(), v.rend());
  int ans = 0;
  std::vector<int> usedc(N, 0);
  stack<int> sta;
  int res = 0;
  int color = 0;
  for (int i = 0; i < K; i++) {
    res += v[i].first;
    if (usedc[v[i].second]++ == 0)
      color++;
    else
      sta.push(v[i].first);
  }
  ans = res + color * color;
  for (int i = K; i < N && sta.size(); i++) {
    if (usedc[v[i].second]++)
      continue;
    int a = sta.top();
    sta.pop();
    res -= a;
    res += v[i].first;
    color++;
    ans = max(ans, res + color * color);
  }
  cout << ans << endl;
}
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
from collections import defaultdict
from heapq import heappush, heappop


n, k = map(int, input().split())
dic = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    dic[t].append(d)

maxs = []
for key in dic:
    dic[key].sort()
    maxs.append((dic[key].pop(), key))

maxs.sort(reverse=True)
queSize = 0
que = []
queScore = 0
defaultScore = 0
cnt = 0
ans = 0
for v, key in maxs:
    defaultScore += v + cnt * 2 + 1
    cnt += 1
    for add in dic[key]:
        heappush(que, add)
        queSize += 1
        queScore += add
    while queSize > k - cnt:
        queScore -= heappop(que)
        queSize -= 1
    ans = max(ans, defaultScore + queScore)

print(ans)


Fixed submission:
from collections import defaultdict
from heapq import heappush, heappop


n, k = map(int, input().split())
dic = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    dic[t].append(d)

maxs = []
for key in dic:
    dic[key].sort()
    maxs.append((dic[key].pop(), key))

maxs.sort(reverse=True)
queSize = 0
que = []
queScore = 0
defaultScore = 0
cnt = 0
ans = 0
for v, key in maxs:
    defaultScore += v + cnt * 2 + 1
    cnt += 1
    if cnt >= k:
        ans = max(ans, defaultScore)
        break
    for add in dic[key]:
        heappush(que, add)
        queSize += 1
        queScore += add
    while queSize > k - cnt:
        queScore -= heappop(que)
        queSize -= 1
    ans = max(ans, defaultScore + queScore)

print(ans)
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
from collections import deque


def main():
    N, K = map(int, input().split())
    sushis = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        sushis[i] = (t, d)
    sushis.sort(key=lambda a: (-a[1], a[0]))
    tmp = 0
    types = set()
    stack = deque()
    ans = 0
    for i in range(len(sushis)):
        t, d = sushis[i]
        if i < K:
            if t in types:
                stack.append(d)
            types.add(t)
            tmp += d
        else:
            if t not in types:
                types.add(t)
                tmp += d
                tmp -= stack.pop()
        ans = max(ans, tmp + len(types) ** 2)
    print(ans)


if __name__ == "__main__":
    main()


Fixed submission:
from collections import deque


def main():
    N, K = map(int, input().split())
    sushis = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        sushis[i] = (t, d)
    sushis.sort(key=lambda a: (-a[1], a[0]))
    tmp = 0
    types = set()
    stack = deque()
    ans = 0
    for i in range(len(sushis)):
        t, d = sushis[i]
        if i < K:
            if t in types:
                stack.append(d)
            types.add(t)
            tmp += d
        else:
            if t not in types and len(stack):
                types.add(t)
                tmp += d
                tmp -= stack.pop()
        ans = max(ans, tmp + len(types) ** 2)
    print(ans)


if __name__ == "__main__":
    main()
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
import heapq

INF = float("inf")
n, k = map(int, input().split())
sushi = {}

for _ in range(n):
    t, d = map(int, input().split())
    if t not in sushi:
        sushi[t] = []
    sushi[t].append(d)

for s in sushi.values():
    s.sort(reverse=True)

eat_counts = {s: 0 for s in sushi.keys()}

# まだ食べてないやつ
news = []
# リピート
repeats = []
# 食べた履歴
eats = []

# 各寿司の一番おいしいやつを heap へ
for name, deli_li in sushi.items():
    # (-おいしさ, name)
    heapq.heappush(news, (-deli_li[0], name))

sum_deli = 0
variety = 0

for _ in range(k):
    # おいしいやつから食べる
    if news and (not repeats or news[0] <= repeats[0]):
        deli, name = heapq.heappop(news)
        variety += 1
    else:
        deli, name = heapq.heappop(repeats)

    sum_deli += -deli
    eat_counts[name] += 1
    heapq.heappush(eats, (-deli, name))
    if eat_counts[name] < len(sushi[name]):
        heapq.heappush(repeats, (-sushi[name][eat_counts[name]], name))

# 美味しいやつから順に食べた
# 種類増やしたらもっと得点上がるか確かめる
ans = sum_deli + variety**2
while True:
    if not news:
        break

    # 食べたやつポイント小さいやつから置き換えたい
    ate_min_deli, ate_min_name = heapq.heappop(eats)
    eat_counts[ate_min_name] -= 1
    if eat_counts[ate_min_name] <= 0:
        # これは種類減っちゃうからだめ
        continue

    sum_deli -= ate_min_deli

    deli, name = heapq.heappop(news)
    variety += 1
    sum_deli += -deli
    eat_counts[name] += 1

    ans = max(ans, sum_deli + variety**2)

print(ans)


Fixed submission:
import heapq

INF = float("inf")
n, k = map(int, input().split())
sushi = {}

for _ in range(n):
    t, d = map(int, input().split())
    if t not in sushi:
        sushi[t] = []
    sushi[t].append(d)

for s in sushi.values():
    s.sort(reverse=True)

eat_counts = {s: 0 for s in sushi.keys()}

# まだ食べてないやつ
news = []
# リピート
repeats = []
# 食べた履歴
eats = []

# 各寿司の一番おいしいやつを heap へ
for name, deli_li in sushi.items():
    # (-おいしさ, name)
    heapq.heappush(news, (-deli_li[0], name))

sum_deli = 0
variety = 0

for _ in range(k):
    # おいしいやつから食べる
    if news and (not repeats or news[0] <= repeats[0]):
        deli, name = heapq.heappop(news)
        variety += 1
    else:
        deli, name = heapq.heappop(repeats)

    sum_deli += -deli
    eat_counts[name] += 1
    heapq.heappush(eats, (-deli, name))
    if eat_counts[name] < len(sushi[name]):
        heapq.heappush(repeats, (-sushi[name][eat_counts[name]], name))

# 美味しいやつから順に食べた
# 種類増やしたらもっと得点上がるか確かめる
ans = sum_deli + variety**2
while True:
    if not news or not eats:
        break

    # 食べたやつポイント小さいやつから置き換えたい
    ate_min_deli, ate_min_name = heapq.heappop(eats)
    eat_counts[ate_min_name] -= 1
    if eat_counts[ate_min_name] <= 0:
        # これは種類減っちゃうからだめ
        continue

    sum_deli -= ate_min_deli

    deli, name = heapq.heappop(news)
    variety += 1
    sum_deli += -deli
    eat_counts[name] += 1

    ans = max(ans, sum_deli + variety**2)

print(ans)
Problem p03148 (Python)

Buggy submission (status: Runtime Error):
N, K = map(int, input().split())
umai = {}
mazui = []
for _ in range(N):
    t, d = map(int, input().split())
    if t not in umai:
        umai[t] = d
    else:
        tmp = umai[t]
        if tmp < d:
            umai[t] = d
            mazui.append((t, tmp))
        else:
            mazui.append((t, d))

umai = sorted(umai.items(), key=lambda x: -x[1])
mazui = sorted(mazui, key=lambda x: -x[1])


umai_len = len(umai)
mazui_len = len(mazui)

acc_umai = [0 for _ in range(umai_len)]
acc_umai[0] = umai[0][1]
for i in range(umai_len - 1):
    acc_umai[i + 1] = acc_umai[i] + umai[i + 1][1]

acc_mazui = [0 for _ in range(mazui_len)]
acc_mazui[0] = mazui[0][1]
for i in range(mazui_len - 1):
    acc_mazui[i + 1] = acc_mazui[i] + mazui[i + 1][1]
acc_mazui.insert(0, 0)

ans = []
for i in range(1, K + 1):
    if umai_len < i or mazui_len < K - i:
        continue
    ans.append(acc_umai[i - 1] + acc_mazui[K - i] + i**2)
print(max(ans))


Fixed submission:
N, K = map(int, input().split())
umai = {}
mazui = []
for _ in range(N):
    t, d = map(int, input().split())
    if t not in umai:
        umai[t] = d
    else:
        tmp = umai[t]
        if tmp < d:
            umai[t] = d
            mazui.append((t, tmp))
        else:
            mazui.append((t, d))

umai = sorted(umai.items(), key=lambda x: -x[1])
mazui = sorted(mazui, key=lambda x: -x[1])


umai_len = len(umai)
mazui_len = len(mazui)

acc_umai = [0 for _ in range(umai_len)]
acc_umai[0] = umai[0][1]
for i in range(umai_len - 1):
    acc_umai[i + 1] = acc_umai[i] + umai[i + 1][1]

acc_mazui = [0 for _ in range(mazui_len)]
if mazui_len > 0:
    acc_mazui[0] = mazui[0][1]
    for i in range(mazui_len - 1):
        acc_mazui[i + 1] = acc_mazui[i] + mazui[i + 1][1]
acc_mazui.insert(0, 0)

ans = []
for i in range(1, K + 1):
    if umai_len < i or mazui_len < K - i:
        continue
    ans.append(acc_umai[i - 1] + acc_mazui[K - i] + i**2)
print(max(ans))
Problem p03148 (Python)

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

    first = set()
    second = []
    point = 0
    for s in sushi[:k]:
        if s[0] not in first:
            first.add(s[0])
        else:
            second.append(s)
        point += s[1]
    point += len(first) ** 2
    ans = point

    for s in sushi[k:]:
        if s[0] not in first:
            point = (
                point + s[1] - second.pop()[1] + (len(first) + 1) ** 2 - len(first) ** 2
            )
            first.add(s[0])
            ans = max(ans, point)
    print(ans)


main()


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

    first = set()
    second = []
    point = 0
    for s in sushi[:k]:
        if s[0] not in first:
            first.add(s[0])
        else:
            second.append(s)
        point += s[1]
    point += len(first) ** 2
    ans = point

    for s in sushi[k:]:
        if len(second) == 0:
            break
        if s[0] not in first:
            point = (
                point + s[1] - second.pop()[1] + (len(first) + 1) ** 2 - len(first) ** 2
            )
            first.add(s[0])
            ans = max(ans, point)
    print(ans)


main()
Problem p03148 (C++)

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

int main() {
  int n, k;
  cin >> n >> k;
  vector<vector<int>> v(n);
  for (int i = 0; i < n; i++) {
    int t, d;
    cin >> t >> d;
    v[t - 1].push_back(d);
  }
  vector<ll> a{0}, b{0};
  for (auto &e : v) {
    if (e.empty()) {
      continue;
    }
    sort(begin(e), end(e));
    a.push_back(e.back());
    for (auto it = begin(e); it != end(e) - 1; it++) {
      b.push_back(*it);
    }
  }
  sort(rbegin(a), rend(a) - 1);
  sort(rbegin(b), rend(b) - 1);
  for (int i = 0; i + 1 < a.size(); i++) {
    a[i + 1] += a[i];
  }
  for (int i = 0; i + 1 < b.size(); i++) {
    b[i + 1] += b[i];
  }
  ll r = 0;
  for (ll x = 1; x <= min(k, (int)a.size()); x++) {
    r = max(r, a[x] + b[k - x] + x * x);
  }
  cout << r << endl;
}


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

int main() {
  int n, k;
  cin >> n >> k;
  vector<vector<int>> v(n);
  for (int i = 0; i < n; i++) {
    int t, d;
    cin >> t >> d;
    v[t - 1].push_back(d);
  }
  vector<ll> a{0}, b{0};
  for (auto &e : v) {
    if (e.empty()) {
      continue;
    }
    sort(begin(e), end(e));
    a.push_back(e.back());
    for (auto it = begin(e); it != end(e) - 1; it++) {
      b.push_back(*it);
    }
  }
  sort(rbegin(a), rend(a) - 1);
  sort(rbegin(b), rend(b) - 1);
  for (int i = 0; i + 1 < a.size(); i++) {
    a[i + 1] += a[i];
  }
  for (int i = 0; i + 1 < b.size(); i++) {
    b[i + 1] += b[i];
  }
  ll r = 0;
  for (ll x = max(1, k - ((int)b.size() - 1)); x <= min(k, (int)a.size() - 1);
       x++) {
    r = max(r, a[x] + b[k - x] + x * x);
  }
  cout << r << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <assert.h>
#include <cstring>
#include <iterator>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define INF 2140000000
const long long MOD = 1000000007;

using namespace std;

int main(int argc, char *argv[]) {
  int n, k;
  scanf("%d%d", &n, &k);
  vector<pair<int, int>> z;
  int i;
  for (i = 0; i < n; i++) {
    int t, d;
    scanf("%d%d", &t, &d);
    z.push_back(make_pair(d, t));
  }
  sort(z.rbegin(), z.rend());

  vector<pair<int, int>> z0;
  set<int> kind;
  ll pts = 0;
  for (i = 0; i < k; i++) {
    pts += z[i].first;
    if (!kind.insert(z[i].second).second) {
      z0.push_back(z[i]); // can be deleted
    }
  }
  reverse(z0.begin(), z0.end());

  ll val = pts + (ll)kind.size() * kind.size();
  ll max = val;

  int id = 0;
  for (i = k; i < n; i++) {
    if (kind.insert(z[i].second).second) {
      pts += (z[i].first - z0[id].first);
      id++;
      val = pts + (ll)kind.size() * kind.size();
      max = MAX(max, val);
    }
    if (id == z0.size()) {
      break;
    }
  }
  printf("%lld\n", max);

  return 0;
}


Fixed submission:
#include <algorithm>
#include <assert.h>
#include <cstring>
#include <iterator>
#include <list>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define INF 2140000000
const long long MOD = 1000000007;

using namespace std;

int main(int argc, char *argv[]) {
  int n, k;
  scanf("%d%d", &n, &k);
  vector<pair<int, int>> z;
  int i;
  for (i = 0; i < n; i++) {
    int t, d;
    scanf("%d%d", &t, &d);
    z.push_back(make_pair(d, t));
  }
  sort(z.rbegin(), z.rend());

  vector<pair<int, int>> z0;
  set<int> kind;
  ll pts = 0;
  for (i = 0; i < k; i++) {
    pts += z[i].first;
    if (!kind.insert(z[i].second).second) {
      z0.push_back(z[i]); // can be deleted
    }
  }
  reverse(z0.begin(), z0.end());

  ll val = pts + (ll)kind.size() * kind.size();
  ll max = val;
  if (z0.empty()) {
    printf("%lld\n", max);
    return 0;
  }

  int id = 0;
  for (i = k; i < n; i++) {
    if (kind.insert(z[i].second).second) {
      pts += (z[i].first - z0[id].first);
      id++;
      val = pts + (ll)kind.size() * kind.size();
      max = MAX(max, val);
    }
    if (id == z0.size()) {
      break;
    }
  }
  printf("%lld\n", max);

  return 0;
}
Problem p03148 (C++)

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

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

  vector<pair<ll, int>> A(N);
  for (int i = 0; i < N; i++) {
    A[i] = make_pair(d[i], t[i]);
  }
  sort(A.begin(), A.end());
  reverse(A.begin(), A.end());
  ll x = 0, y = 0;
  vector<bool> used(N);
  priority_queue<ll, vector<ll>, greater<ll>> que;
  for (int i = 0; i < K; i++) {
    y += A[i].first;
    if (!used[A[i].second]) {
      x++;
      used[A[i].second] = true;
    } else {
      que.push(A[i].first);
    }
  }
  ll ans = y + x * x;
  for (int i = K; i < N; i++) {
    if (used[A[i].second])
      continue;
    y -= que.top();
    que.pop();
    y += A[i].first;
    x++;
    used[A[i].second] = true;
    ans = max(ans, y + x * x);
  }
  cout << ans << endl;
}

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

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

  vector<pair<ll, int>> A(N);
  for (int i = 0; i < N; i++) {
    A[i] = make_pair(d[i], t[i]);
  }
  sort(A.begin(), A.end());
  reverse(A.begin(), A.end());
  ll x = 0, y = 0;
  vector<bool> used(N);
  priority_queue<ll, vector<ll>, greater<ll>> que;
  for (int i = 0; i < K; i++) {
    y += A[i].first;
    if (!used[A[i].second]) {
      x++;
      used[A[i].second] = true;
    } else {
      que.push(A[i].first);
    }
  }
  ll ans = y + x * x;
  for (int i = K; i < N && !que.empty(); i++) {
    if (used[A[i].second])
      continue;
    y -= que.top();
    que.pop();
    y += A[i].first;
    x++;
    used[A[i].second] = true;
    ans = max(ans, y + x * x);
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <functional>
#include <queue>
#include <stdio.h>
#include <vector>
using namespace std;
using ll = long long;

using pii = pair<ll, ll>;
ll n, k;
pii s[100000];
vector<ll> u, v;
int main() {
  scanf("%lld %lld", &n, &k);
  for (int i = 0; i < n; i++)
    scanf("%lld %lld", &s[i].first, &s[i].second);
  sort(s, s + n);
  for (int i = 0; i < n; i++) {
    if (n <= i + 1 || s[i].first != s[i + 1].first)
      u.emplace_back(s[i].second);
    else
      v.emplace_back(s[i].second);
  }
  sort(u.begin(), u.end(), greater<ll>());
  sort(v.begin(), v.end(), greater<ll>());
  priority_queue<ll> hoketu;
  while (k < u.size() + v.size()) {
    hoketu.emplace(v.back());
    v.pop_back();
  }

  ll max_r = 0, cur = 0;
  for (ll e : u)
    cur += e;
  for (ll e : v)
    cur += e;
  while (!u.empty()) {
    max_r = max(max_r, cur + ((ll)u.size()) * ((ll)u.size()));
    cur -= u.back();
    hoketu.emplace(u.back());
    u.pop_back();

    cur += hoketu.top();
    hoketu.pop();
  }
  printf("%lld\n", max(max_r, cur));
}

Fixed submission:
#include <algorithm>
#include <functional>
#include <queue>
#include <stdio.h>
#include <vector>
using namespace std;
using ll = long long;

using pii = pair<ll, ll>;
ll n, k;
pii s[100000];
vector<ll> u, v;
int main() {
  scanf("%lld %lld", &n, &k);
  for (int i = 0; i < n; i++)
    scanf("%lld %lld", &s[i].first, &s[i].second);
  sort(s, s + n);
  for (int i = 0; i < n; i++) {
    if (n <= i + 1 || s[i].first != s[i + 1].first)
      u.emplace_back(s[i].second);
    else
      v.emplace_back(s[i].second);
  }
  sort(u.begin(), u.end(), greater<ll>());
  sort(v.begin(), v.end(), greater<ll>());
  priority_queue<ll> hoketu;
  while (k < u.size() + v.size()) {
    if (v.empty()) {
      hoketu.emplace(u.back());
      u.pop_back();
    } else {
      hoketu.emplace(v.back());
      v.pop_back();
    }
  }

  ll max_r = 0, cur = 0;
  for (ll e : u)
    cur += e;
  for (ll e : v)
    cur += e;
  while (!u.empty()) {
    max_r = max(max_r, cur + ((ll)u.size()) * ((ll)u.size()));
    cur -= u.back();
    hoketu.emplace(u.back());
    u.pop_back();

    cur += hoketu.top();
    hoketu.pop();
  }
  printf("%lld\n", max(max_r, cur));
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <assert.h>
#include <cmath>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
using namespace std;

using int64 = long long;
using int128 = __int128;

#define repeat(i, n) for (auto i = decltype(n)(); (i) < (n); (i)++)

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

  int N, K;
  cin >> N >> K;

  vector<int64> t(N), d(N);
  repeat(i, N) { cin >> t[i] >> d[i]; }
  vector<int> indices(N);
  iota(begin(indices), end(indices), 0);
  sort(begin(indices), end(indices), [&](int i, int j) { return d[i] > d[j]; });

  using P = pair<int, int>; // (美味しさ, 種類)

  int64 x = 0, sum_d = 0;
  vector<int> counter(N, 0);
  priority_queue<P, vector<P>, greater<P>> pque;

  repeat(i, K) {
    int idx = indices[i];
    sum_d += d[idx];
    counter[t[idx]]++;

    if (counter[t[idx]] == 1)
      x++;
    else
      pque.push(P(d[idx], t[idx]));
  }

  int64 ans = sum_d + x * x;

  for (int i = K; i < N; i++) {
    int idx = indices[i];

    if (!pque.empty() and counter[t[idx]] == 0) {
      sum_d -= pque.top().first;
      pque.pop();
      sum_d += d[idx];
      x++;
      counter[t[idx]]++;
    }
    ans = max(ans, sum_d + x * x);
  }
  cout << ans << endl;

  return 0;
}


Fixed submission:
#include <algorithm>
#include <assert.h>
#include <cmath>
#include <cstring>
#include <deque>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
using namespace std;

using int64 = long long;
using int128 = __int128;

#define repeat(i, n) for (auto i = decltype(n)(); (i) < (n); (i)++)

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

  int N, K;
  cin >> N >> K;

  vector<int64> t(N), d(N);
  repeat(i, N) {
    cin >> t[i] >> d[i];
    t[i]--;
  }
  vector<int> indices(N);
  iota(begin(indices), end(indices), 0);
  sort(begin(indices), end(indices), [&](int i, int j) { return d[i] > d[j]; });

  using P = pair<int, int>; // (美味しさ, 種類)

  int64 x = 0, sum_d = 0;
  vector<int> counter(N, 0);
  priority_queue<P, vector<P>, greater<P>> pque;

  repeat(i, K) {
    int idx = indices[i];
    sum_d += d[idx];
    counter[t[idx]]++;

    if (counter[t[idx]] == 1)
      x++;
    else
      pque.push(P(d[idx], t[idx]));
  }

  int64 ans = sum_d + x * x;

  for (int i = K; i < N; i++) {
    int idx = indices[i];

    if (!pque.empty() and counter[t[idx]] == 0) {
      sum_d -= pque.top().first;
      pque.pop();
      sum_d += d[idx];
      x++;
      counter[t[idx]]++;
    }
    ans = max(ans, sum_d + x * x);
  }
  cout << ans << endl;

  return 0;
}
Problem p03148 (C++)

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

int count_and_compress(vector<long long> &a) {
  vector<long long> b = a;
  sort(b.begin(), b.end());
  b.erase(unique(b.begin(), b.end()));
  for (auto &e : a) {
    e = (lower_bound(b.begin(), b.end(), e) - b.begin());
  }
  return b.size();
}

int main() {
  int n, k;
  cin >> n >> k;
  vector<long long> t(n), d(n);
  for (int i = 0; i < n; i++) {
    cin >> t[i] >> d[i];
  }

  int sz = count_and_compress(t);
  vector<vector<long long>> g(sz);
  for (int i = 0; i < n; i++) {
    g[t[i]].emplace_back(d[i]);
  }
  for (auto &v : g) {
    sort(v.begin(), v.end(), greater<long long>());
  }

  vector<long long> f(1, 0), s(1, 0);
  for (int i = 0; i < sz; i++) {
    for (int j = 0; j < (int)g[i].size(); j++) {
      if (j == 0)
        f.emplace_back(g[i][j]);
      else
        s.emplace_back(g[i][j]);
    }
  }

  sort(f.begin() + 1, f.end(), greater<long long>());
  sort(s.begin() + 1, s.end(), greater<long long>());
  for (long long i = 1; i < (int)f.size(); i++) {
    f[i] += (2 * i - 1);
  }

  for (int i = 0; i + 1 < (int)f.size(); i++) {
    f[i + 1] += f[i];
  }
  for (int i = 0; i + 1 < (int)s.size(); i++) {
    s[i + 1] += s[i];
  }

  long long ans = 0;
  for (int i = 1; i <= k; i++) {
    int j = k - i;
    if (i < f.size() and j < s.size()) {
      long long res = f[i] + s[j];
      ans = max(ans, res);
    }
  }

  cout << ans << endl;
  return 0;
}

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

int count_and_compress(vector<long long> &a) {
  vector<long long> b = a;
  sort(b.begin(), b.end());
  b.erase(unique(b.begin(), b.end()), b.end());
  for (auto &e : a) {
    e = (lower_bound(b.begin(), b.end(), e) - b.begin());
  }
  return b.size();
}

int main() {
  int n, k;
  cin >> n >> k;
  vector<long long> t(n), d(n);
  for (int i = 0; i < n; i++) {
    cin >> t[i] >> d[i];
  }

  int sz = count_and_compress(t);
  vector<vector<long long>> g(sz);
  for (int i = 0; i < n; i++) {
    g[t[i]].emplace_back(d[i]);
  }
  for (auto &v : g) {
    sort(v.begin(), v.end(), greater<long long>());
  }

  vector<long long> f(1, 0), s(1, 0);
  for (int i = 0; i < sz; i++) {
    for (int j = 0; j < (int)g[i].size(); j++) {
      if (j == 0)
        f.emplace_back(g[i][j]);
      else
        s.emplace_back(g[i][j]);
    }
  }

  sort(f.begin() + 1, f.end(), greater<long long>());
  sort(s.begin() + 1, s.end(), greater<long long>());
  for (long long i = 1; i < (int)f.size(); i++) {
    f[i] += (2 * i - 1);
  }

  for (int i = 0; i + 1 < (int)f.size(); i++) {
    f[i + 1] += f[i];
  }
  for (int i = 0; i + 1 < (int)s.size(); i++) {
    s[i + 1] += s[i];
  }

  long long ans = 0;
  for (int i = 1; i <= k; i++) {
    int j = k - i;
    if (i < f.size() and j < s.size()) {
      long long res = f[i] + s[j];
      ans = max(ans, res);
    }
  }

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

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

struct ST {
  ll t, d;
};
bool operator<(const ST &a, const ST &b) { return a.d > b.d; }

int main() {
  int N, K;
  cin >> N >> K;
  vector<ST> v;
  rep(i, N) {
    ll t, d;
    cin >> t >> d;
    v.push_back((ST){t, d});
  }
  sort(ALLOF(v));

  ll sum = 0;
  map<ll, ll> mx, nx;
  vector<ll> w;
  vector<ST> w2;
  rep(i, N) {
    if (i < K) {
      sum += v[i].d;
      if (mx.count(v[i].t) > 0)
        w.push_back(v[i].d);
      else {
        mx[v[i].t] = v[i].d;
      }
    } else {
      if (mx.count(v[i].t) == 0 && nx.count(v[i].t) == 0) {
        nx[v[i].t] = v[i].d;
        w2.push_back(v[i]);
      }
    }
  }

  ll x = mx.size();
  ll ret = sum + x * x;
  rep(i, w2.size()) {
    if (w.size() - 1 - i < 0)
      break;
    sum -= w[w.size() - 1 - i];
    sum += w2[i].d;
    x++;
    ret = max(ret, sum + x * x);
  }

  cout << ret << endl;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, n) for (int i = (a); i < (int)(n); i++)
#define rep(i, n) REP(i, 0, n)
#define FOR(it, c)                                                             \
  for (__typeof((c).begin()) it = (c).begin(); it != (c).end(); ++it)
#define ALLOF(c) (c).begin(), (c).end()
typedef long long ll;
typedef unsigned long long ull;

struct ST {
  ll t, d;
};
bool operator<(const ST &a, const ST &b) { return a.d > b.d; }

int main() {
  int N, K;
  cin >> N >> K;
  vector<ST> v;
  rep(i, N) {
    ll t, d;
    cin >> t >> d;
    v.push_back((ST){t, d});
  }
  sort(ALLOF(v));

  ll sum = 0;
  map<ll, ll> mx, nx;
  vector<ll> w;
  vector<ST> w2;
  rep(i, N) {
    if (i < K) {
      sum += v[i].d;
      if (mx.count(v[i].t) > 0)
        w.push_back(v[i].d);
      else {
        mx[v[i].t] = v[i].d;
      }
    } else {
      if (mx.count(v[i].t) == 0 && nx.count(v[i].t) == 0) {
        nx[v[i].t] = v[i].d;
        w2.push_back(v[i]);
      }
    }
  }

  ll x = mx.size();
  ll ret = sum + x * x;
  rep(i, w2.size()) {
    if (x >= K)
      break;
    if (w.size() - 1 - i < 0)
      break;
    sum -= w[w.size() - 1 - i];
    sum += w2[i].d;
    x++;
    ret = max(ret, sum + x * x);
  }

  cout << ret << endl;

  return 0;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
#define rep(i, n) for (int i = 0; i < n; i++)
#define repb(i, a, b) for (int i = a; i < b; i++)

using namespace std;
typedef pair<int, int> P;
typedef long long int ll;

const int INF = 1 << 30;

int main() {
  int n, k;
  cin >> n >> k;
  vector<P> sushi(n);
  rep(i, n) {
    int t, d;
    cin >> t >> d;
    sushi[i] = P(d, --t);
  }
  sort(sushi.begin(), sushi.end(), greater<P>());
  vector<bool> include(k, false);
  vector<int> selected;
  ll kind = 0LL;
  ll sum = 0LL;
  rep(i, k) {
    P p = sushi[i];
    if (!include[p.second]) {
      include[p.second] = true;
      kind++;
    } else {
      selected.push_back(p.first);
    }
    sum += p.first;
  }
  sort(selected.begin(), selected.end(), greater<int>());
  ll ans = sum + kind * kind;
  ll id = (ll)kind;
  while (!selected.empty() && id < n) {
    if (include[sushi[id].second]) {
      id++;
      continue;
    }
    include[sushi[id].second] = true;
    sum -= selected.back();
    selected.pop_back();
    sum += sushi[id].first;
    id++;
    kind++;
    ans = max(ans, sum + kind * kind);
  }
  cout << ans << endl;
  // system("read -n 1 -s -p \"Press any key to continue...\"");
}


Fixed submission:
#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
#define rep(i, n) for (int i = 0; i < n; i++)
#define repb(i, a, b) for (int i = a; i < b; i++)

using namespace std;
typedef pair<int, int> P;
typedef long long int ll;

const int INF = 1 << 30;

int main() {
  int n, k;
  cin >> n >> k;
  vector<P> sushi(n);
  rep(i, n) {
    int t, d;
    cin >> t >> d;
    sushi[i] = P(d, --t);
  }
  sort(sushi.begin(), sushi.end(), greater<P>());
  vector<bool> include(n, false);
  vector<int> selected;
  ll kind = 0LL;
  ll sum = 0LL;
  rep(i, k) {
    P p = sushi[i];
    if (!include[p.second]) {
      include[p.second] = true;
      kind++;
    } else {
      selected.push_back(p.first);
    }
    sum += p.first;
  }
  sort(selected.begin(), selected.end(), greater<int>());
  ll ans = sum + kind * kind;
  ll id = (ll)kind;
  while (!selected.empty() && id < n) {
    if (include[sushi[id].second]) {
      id++;
      continue;
    }
    include[sushi[id].second] = true;
    sum -= selected.back();
    selected.pop_back();
    sum += sushi[id].first;
    id++;
    kind++;
    ans = max(ans, sum + kind * kind);
  }
  cout << ans << endl;
  // system("read -n 1 -s -p \"Press any key to continue...\"");
}
Problem p03148 (C++)

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

using namespace std;

typedef long long ll;

const ll MOD = 1000000007;
const int MAXN = 100000;

vector<ll> sushi[MAXN + 1];
vector<ll> best;
vector<ll> rest;

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  ll n, k;
  cin >> n >> k;
  ll t[n], d[n];
  for (int i = 0; i < n; i++)
    cin >> t[i] >> d[i];
  for (int i = 0; i < n; i++) {
    sushi[t[i]].push_back(d[i]);
  }
  for (int i = 1; i <= MAXN; i++) {
    if (sushi[i].empty())
      continue;
    sort(sushi[i].rbegin(), sushi[i].rend());
    best.push_back(sushi[i][0]);
    for (int j = 1; j < sushi[i].size(); j++) {
      rest.push_back(sushi[i][j]);
    }
  }
  sort(best.rbegin(), best.rend());
  sort(rest.rbegin(), rest.rend());
  ll bsize = best.size();
  ll rsize = rest.size();
  ll sumbest[bsize], sumrest[rsize];
  sumbest[0] = best[0];
  for (int i = 1; i < bsize; i++)
    sumbest[i] = sumbest[i - 1] + best[i];
  sumrest[0] = rest[0];
  for (int i = 1; i < rsize; i++)
    sumrest[i] = sumrest[i - 1] + rest[i];
  ll ans = 0;
  for (ll i = 0; i <= min(k, bsize); i++) {
    // best * i & rest * (k - i)
    ll sum = 0;
    if (i - 1 >= 0) {
      sum += sumbest[i - 1];
    }
    if (k - i - 1 >= 0 && k - i - 1 < rsize) {
      sum += sumrest[k - i - 1];
    }
    ans = max(ans, sum + i * i);
  }
  cout << ans << endl;
}

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

using namespace std;

typedef long long ll;

const ll MOD = 1000000007;
const int MAXN = 100000;

vector<ll> sushi[MAXN + 1];
vector<ll> best;
vector<ll> rest;

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  ll n, k;
  cin >> n >> k;
  ll t[n], d[n];
  for (int i = 0; i < n; i++)
    cin >> t[i] >> d[i];
  for (int i = 0; i < n; i++) {
    sushi[t[i]].push_back(d[i]);
  }
  for (int i = 1; i <= MAXN; i++) {
    if (sushi[i].empty())
      continue;
    sort(sushi[i].rbegin(), sushi[i].rend());
    best.push_back(sushi[i][0]);
    for (int j = 1; j < sushi[i].size(); j++) {
      rest.push_back(sushi[i][j]);
    }
  }
  sort(best.rbegin(), best.rend());
  sort(rest.rbegin(), rest.rend());
  ll bsize = best.size();
  ll rsize = rest.size();
  ll sumbest[bsize], sumrest[rsize];
  sumbest[0] = best[0];
  for (int i = 1; i < bsize; i++)
    sumbest[i] = sumbest[i - 1] + best[i];
  if (!rest.empty()) {
    sumrest[0] = rest[0];
  }
  for (int i = 1; i < rsize; i++)
    sumrest[i] = sumrest[i - 1] + rest[i];
  ll ans = 0;
  for (ll i = 0; i <= min(k, bsize); i++) {
    // best * i & rest * (k - i)
    ll sum = 0;
    if (i - 1 >= 0) {
      sum += sumbest[i - 1];
    }
    if (k - i - 1 >= 0 && k - i - 1 < rsize) {
      sum += sumrest[k - i - 1];
    }
    ans = max(ans, sum + i * i);
  }
  cout << ans << endl;
}
Problem p03148 (C++)

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

using namespace std;

typedef pair<int, int> P;

int main() {
  int N, K;
  cin >> N >> K;
  P p[100000];
  long s = 0;
  int cnt[100001] = {0};
  long ans = 0;
  for (int i = 0; i < N; i++) {
    int t, d;
    cin >> t >> d;
    p[i] = P(d, t);
  }
  sort(p, p + N, greater<P>());
  for (int i = 0; i < K; i++) {
    if (cnt[p[i].second] == 0)
      s++;
    cnt[p[i].second]++;
    ans += p[i].first;
  }
  ans += s * s;
  long ans_m = ans;
  int l = K - 1;
  int r = K;
  while (l >= 0 && K <= N - 1) {
    if (/*2*s+1 >= (p[l].first - p[r].first) && */ cnt[p[l].second] >= 2 &&
        cnt[p[r].second] == 0) {
      ans += 2 * s + 1 - (p[l].first - p[r].first);
      cnt[p[l].second]--;
      cnt[p[r].second]++;
      l--;
      r++;
      s++;
    } else if (cnt[p[l].second] <= 1) {
      l--;
    } else if (cnt[p[r].second] > 0) {
      r++;
    } /*else if(2*s+1 < (p[l].first - p[r].first)){
         //l--;
     }*/
    ans_m = max(ans, ans_m);
  }
  cout << ans_m << endl;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <queue>

using namespace std;

typedef pair<int, int> P;

int main() {
  int N, K;
  cin >> N >> K;
  P p[100000];
  long s = 0;
  int cnt[100001] = {0};
  long ans = 0;
  for (int i = 0; i < N; i++) {
    int t, d;
    cin >> t >> d;
    p[i] = P(d, t);
  }
  sort(p, p + N, greater<P>());
  for (int i = 0; i < K; i++) {
    if (cnt[p[i].second] == 0)
      s++;
    cnt[p[i].second]++;
    ans += p[i].first;
  }
  ans += s * s;
  long ans_m = ans;
  int l = K - 1;
  int r = K;
  while (l >= 0 && r <= N - 1) {
    if (/*2*s+1 >= (p[l].first - p[r].first) && */ cnt[p[l].second] >= 2 &&
        cnt[p[r].second] == 0) {
      ans += 2 * s + 1 - (p[l].first - p[r].first);
      cnt[p[l].second]--;
      cnt[p[r].second]++;
      l--;
      r++;
      s++;
    } else if (cnt[p[l].second] <= 1) {
      l--;
    } else if (cnt[p[r].second] > 0) {
      r++;
    } /*else if(2*s+1 < (p[l].first - p[r].first)){
         //l--;
     }*/
    ans_m = max(ans, ans_m);
  }
  cout << ans_m << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
vector<int> v[100001];
int n, k;
int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    ll t, d;
    cin >> t >> d;
    v[t].push_back(d);
  }
  // 種類数が増えない
  vector<ll> ans0;
  // 種類数が増える
  vector<ll> ans1;
  for (int i = 0; i <= n; i++) {
    if (v[i].empty()) {
      continue;
    }
    sort(v[i].begin(), v[i].end(), greater<ll>());
    ans1.push_back(v[i][0]);
    for (int j = 1; j < v[i].size(); j++) {
      ans0.push_back(v[i][j]);
    }
  }
  sort(ans0.begin(), ans0.end(), greater<ll>());
  sort(ans1.begin(), ans1.end(), greater<ll>());
  ll ans = 0LL;
  for (ll i = 1; i <= k; i++) {
    if (ans0.size() < k - i) {
      continue;
    }
    if (ans1.size() < i) {
      continue;
    }
    ll t = accumulate(ans1.begin(), ans1.begin() + i, 0LL) +
           accumulate(ans0.begin(), ans0.begin() + k - i, 0LL) + i * i;
    ans = max(ans, t);
  }
  ll t = accumulate(ans1.begin(), ans1.begin() + k, 0LL) + k * k;
  ans = max(ans, t);
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
vector<int> v[100001];
int n, k;
int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    ll t, d;
    cin >> t >> d;
    v[t].push_back(d);
  }
  // 種類数が増えない
  vector<ll> ans0;
  // 種類数が増える
  vector<ll> ans1;
  for (int i = 0; i <= n; i++) {
    if (v[i].empty()) {
      continue;
    }
    sort(v[i].begin(), v[i].end(), greater<ll>());
    ans1.push_back(v[i][0]);
    for (int j = 1; j < v[i].size(); j++) {
      ans0.push_back(v[i][j]);
    }
  }
  sort(ans0.begin(), ans0.end(), greater<ll>());
  sort(ans1.begin(), ans1.end(), greater<ll>());
  ll ans = 0LL;
  for (ll i = 1; i <= k; i++) {
    if (ans0.size() < k - i) {
      continue;
    }
    if (ans1.size() < i) {
      continue;
    }
    ll t = accumulate(ans1.begin(), ans1.begin() + i, 0LL) +
           accumulate(ans0.begin(), ans0.begin() + k - i, 0LL) + i * i;
    ans = max(ans, t);
  }
  cout << ans << endl;
}
Problem p03148 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
using namespace std;
typedef long long int ll;
typedef pair<ll, ll> P;

#define EPS (1e-7)
#define INF 1e18
#define max(p, q) ((p) > (q) ? (p) : (q))
#define min(p, q) ((p) < (q) ? (p) : (q))
#define PI (acos(-1))

#define REP(i, n) for (int i = 0; i < (int)(n); i++)

int main() {

  ll N, K;
  cin >> N >> K;
  vector<P> t(N);
  REP(i, N) cin >> t[i].first >> t[i].second;
  sort(t.begin(), t.end(), [](P a, P b) { return a.second > b.second; });
  priority_queue<ll, vector<ll>, greater<ll>> red;
  set<ll> a;

  ll sum = 0;
  REP(i, K) {
    sum += t[i].second;
    if (a.find(t[i].first) != a.end())
      red.push(t[i].second);
    else
      a.insert(t[i].first);
  }
  ll tmp = sum + a.size() * a.size();
  for (int i = K; i < N; i++) {
    if (a.find(t[i].first) == a.end()) {
      a.insert(t[i].first);
      sum -= red.top();
      sum += t[i].second;
      red.pop();
      tmp = max(tmp, sum + a.size() * a.size());
    }
  }
  cout << tmp << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
using namespace std;
typedef long long int ll;
typedef pair<ll, ll> P;

#define EPS (1e-7)
#define INF 1e18
#define max(p, q) ((p) > (q) ? (p) : (q))
#define min(p, q) ((p) < (q) ? (p) : (q))
#define PI (acos(-1))

#define REP(i, n) for (int i = 0; i < (int)(n); i++)

int main() {

  ll N, K;
  cin >> N >> K;
  vector<P> t(N);
  REP(i, N) cin >> t[i].first >> t[i].second;
  sort(t.begin(), t.end(), [](P a, P b) { return a.second > b.second; });
  priority_queue<ll, vector<ll>, greater<ll>> red;
  set<ll> a;

  ll sum = 0;
  REP(i, K) {
    sum += t[i].second;
    if (a.find(t[i].first) != a.end())
      red.push(t[i].second);
    else
      a.insert(t[i].first);
  }
  ll tmp = sum + a.size() * a.size();
  for (int i = K; i < N; i++) {
    if (a.find(t[i].first) == a.end() && a.size() != K) {
      a.insert(t[i].first);
      sum -= red.top();
      sum += t[i].second;
      red.pop();
      tmp = max(tmp, sum + a.size() * a.size());
    }
  }
  cout << tmp << endl;

  return 0;
}
Problem p03148 (C++)

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

using namespace std;

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

#define MOD 1000000007
#define REP(i, N) for (int i = 0; i < N; ++i)
#define REP1(i, N) for (int i = 1; i <= N; ++i)
#define RREP(i, N) for (int i = N - 1; i >= 0; --i)
#define ALL(a) a.begin(), a.end()

int main() {
  int n, k;
  cin >> n >> k;
  vector<P> vec;
  REP(i, n) {
    int t, d;
    cin >> t >> d;
    t--;
    vec.push_back(P(d, t));
  }
  sort(ALL(vec));
  reverse(ALL(vec));
  set<int> set;
  stack<int> stack;
  ll sum = 0;
  REP(i, k) {
    if (set.count(vec[i].second)) {
      stack.push(vec[i].first);
    } else {
      set.insert(vec[i].second);
    }
    sum += vec[i].first;
  }
  ll ans = sum + set.size() * set.size();
  for (int i = k; i < n; ++i) {
    if (set.count(vec[i].second))
      continue;
    sum -= stack.top();
    stack.pop();
    set.insert(vec[i].second);
    sum += vec[i].first;
    ll res = sum + set.size() * set.size();
    ans = max(ans, res);
  }
  cout << ans << endl;
  return 0;
}

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

using namespace std;

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

#define MOD 1000000007
#define REP(i, N) for (int i = 0; i < N; ++i)
#define REP1(i, N) for (int i = 1; i <= N; ++i)
#define RREP(i, N) for (int i = N - 1; i >= 0; --i)
#define ALL(a) a.begin(), a.end()

int main() {
  int n, k;
  cin >> n >> k;
  vector<P> vec;
  REP(i, n) {
    int t, d;
    cin >> t >> d;
    t--;
    vec.push_back(P(d, t));
  }
  sort(ALL(vec));
  reverse(ALL(vec));
  set<int> set;
  stack<int> stack;
  ll sum = 0;
  REP(i, k) {
    if (set.count(vec[i].second)) {
      stack.push(vec[i].first);
    } else {
      set.insert(vec[i].second);
    }
    sum += vec[i].first;
  }
  ll ans = sum + set.size() * set.size();
  for (int i = k; i < n; ++i) {
    if (set.count(vec[i].second))
      continue;
    if (stack.empty())
      break;
    sum -= stack.top();
    stack.pop();
    set.insert(vec[i].second);
    sum += vec[i].first;
    ll res = sum + set.size() * set.size();
    ans = max(ans, res);
  }
  cout << ans << endl;
  return 0;
}