Problem p03148 (python)

Fixed submission:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi.sort(key=lambda x: x[1], reverse=True)

score = 0

remains = []
kinds = set()
for t, d in sushi[:K]:
    if t in kinds:
        remains.append(d)
    kinds.add(t)
    score += d
score += len(kinds)**2

test = score
for t, d in sushi[K:]:
    if not remains:
        break
    if t in kinds:
        continue
    kinds.add(t)
    rm = remains.pop()
    test = test + (d - rm) + (len(kinds)*2 - 1)
    score = max(score, test)

print(score)

Buggy submission (verdict: Runtime Error):
from heapq import heappop, heappush, heapify

N, K = map(int, input().split())
sushi = [(-d, t) for t, d in map(int, input().split())]
heapify(sushi)

maxs = []
rems = []
kinds = set()
for _ in range(K):
    d, t = heappop(sushi)
    if t in kinds:
        heappush(rems, -d)
    else:
        maxs.append(-d)
        kinds.add(t)
ans = sum(maxs) + sum(rems) + len(kinds)**2

while True:
    if not sushi or not rems:
        break
    d, t = heappop(sushi)
    if t in kinds:
        continue
    else:
        heappop(rems)
        maxs.append(-d)
        kinds.add(t)
        ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Buggy submission (verdict: Runtime Error):
from heapq import heappop, heappush

N, K = map(int, input().split())
sushi = []
for _ in range(N):
    t, d = map(int, input().split())
    heappush(sushi, (-d, t))

maxs = []
rems = []
kinds = set()
for _ in range(K):
    d, t = heappop(sushi)
    if t in kinds:
        heappush(rems, -d)
    else:
        maxs.append(-d)
        kinds.add(t)
ans = sum(maxs) + sum(rems) + len(kinds)**2

while True:
    if not rems:
        break
    d, t = heappop(sushi)
    if t in kinds:
        continue
    else:
        heappop(rems)
        maxs.append(-d)
        kinds.add(t)
        ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Buggy submission (verdict: Runtime Error):
from itertools import groupby
from heapq import heapify, heappop

N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]

srt = sorted(sushi, key=lambda x: x[1], reverse=True)
eat, non_eat = srt[:K], srt[K:]

maxs = []
rems = []
kinds = set()
for k, g in groupby(sorted(eat, key=lambda x: x[0]), key=lambda x: x[0]):
    d = list(list(zip(*g))[1])
    maxs.append(d.pop(0))
    rems.extend(d)
    kinds.add(k)
heapify(rems)
ans = sum(maxs) + sum(rems) + len(kinds)**2

for k, g in groupby(sorted(non_eat, key=lambda x: x[0]), key=lambda x: x[0]):
    if k in kinds:
        continue
    d = list(list(zip(*g))[1])
    maxs.append(d.pop(0))
    heappop(rems)
    kinds.add(k)
    ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Buggy submission (verdict: Runtime Error):
from itertools import groupby
from heapq import heapify, heappop

N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]


def point(sushi):
    t, d = zip(*sushi)
    return sum(d) + len(set(t))**2


srt = sorted(sushi, key=lambda x: x[1], reverse=True)
eat, non_eat = srt[:K], srt[K:]

maxs = []
rems = []
kinds = set()
for k, g in groupby(sorted(eat, key=lambda x: x[0]), key=lambda x: x[0]):
    _, d = zip(*g)
    maxs.append(d.pop(0))
    rems.extend(d)
    kinds.add(k)
heapify(rems)
ans = point(eat)

for k, g in groupby(sorted(non_eat, key=lambda x: x[0]), key=lambda x: x[0]):
    if k in kinds:
        continue
    _, d = zip(*g)
    maxs.append(d.pop(0))
    heappop(rems)
    kinds.add(k)
    ans = max(ans, point(maxs + rems))

print(ans)

Buggy submission (verdict: Runtime Error):
from itertools import groupby

N, K = map(int, input().split())
sushi = sorted([tuple(map(int, input().split())) for _ in range(N)],
               key=lambda s: s[1])

maximums = []
remains = []
v = set()
for k, g in groupby(sushi[-K:], key=lambda x: x[0]):
    v.add(k)
    ds = list(list(zip(*g))[1])
    maximums.append(ds.pop())
    remains.extend(ds)
remains = sorted(remains)
ans = sum(maximums) + sum(remains) + len(v)**2

max_v = max(K, len(set(list(zip(*sushi))[0])))

for k, g in groupby(sorted(sushi[:-K], key=lambda x: x[0]), key=lambda x: x[0]):
    if k in v:
        continue
    v.add(k)
    ds = list(list(zip(*g))[1])
    maximums.append(ds.pop())
    remains.pop()
    ans = max(ans, sum(maximums) + sum(remains) + len(v)**2)

print(ans)