Problem p03148 (python)

Fixed submission:
import heapq

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        heapq.heappush(x, di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
for i in range(n - k - 1, -1, -1):
    if len(x) == 0:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - heapq.heappop(x)
        p += 1
        ans = max(ans, z + p ** 2)

print(ans)

Buggy submission (verdict: Runtime Error):
import heapq

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        heapq.heappush(x, di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

ans = z + len(t.keys()) ** 2
for i in range(n - k - 1, -1, -1):
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - heapq.heappop(x)
        ans = max(ans, z + len(t.keys()) ** 2)
        k += 1

print(ans)

Buggy submission (verdict: Runtime Error):
n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
        t[sushi[i][0]] += 1
    else:
        t[sushi[i][0]] = 1
    z += sushi[i][1]

y = list(t.keys())
ans = z + len(y) ** 2
k = 1
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in y:
        y.append(sushi[i][0])
        xi = x[-k]
        z += sushi[i][1] - sushi[xi][1]
        ans = max(ans, z + len(y) ** 2)
        k += 1

print(ans)

Buggy submission (verdict: Runtime Error):
from collections import defaultdict

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = defaultdict(int)
x = list(range(n - 1, n - k - 1, -1))
z = [0]
for i in range(n - k, n):
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]

y = []
yd = list(t.keys())
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])
y = y[::-1]
k = 1

for k in range(1, len(y) + 1):
    for i in range(len(x) - 1, -1, -1):
        xi = x[i]
        if t[sushi[xi][0]] > 1:
            t[sushi[xi][0]] -= 1
            z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])
            x.pop(i)
            break

p = len(t.keys())
ans = 0
for i in range(len(z)):
    ans = max(z[i] + (p + i) ** 2)
print(ans)

Buggy submission (verdict: Runtime Error):
import copy
from collections import defaultdict

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = [0]
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
        t[sushi[i][0]] += 1
    else:
        t[sushi[i][0]] = 1
    z[0] += sushi[i][1]
y = []
yd = []
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in t.keys() and sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])

for k in range(1, len(x) + 1):
    xi = x[-k]
    z.append(z[-1] - sushi[xi][1] + sushi[y[k - 1]][1])

ans = 0
p = len(t.keys())
for i in range(len(z)):
    ans = max(ans, z[i] + (p + i) ** 2)
print(ans)

Buggy submission (verdict: Runtime Error):
import copy
from collections import defaultdict

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = defaultdict(int)
x = []
z = [0]
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]
y = []
yd = []
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in t.keys() and sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])
y = y[::-1]
k = 1

for k in range(1, len(y) + 1):
    xi = x[-k]
    z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])

p = len(t.keys())
for i in range(len(z)):
    z[i] += (p + i) ** 2
print(max(z))