Problem p03148 (python)

Fixed submission:
from collections import defaultdict as dd
import heapq

n,k=map(int,input().split())
l=[list(map(int,input().split())) for _ in range(n)]
l.sort(key=lambda x:-x[1])

sushi_types = set()
sushies = []
v=0

for i in range(k):
    t,d=l[i]
    v+=d
    if t in sushi_types:
        heapq.heappush(sushies,d)
    else:
        sushi_types.add(t)

remaining_sushi=dd(list)
for i in range(k,n):
    t,d=l[i]
    if not t in sushi_types:
        heapq.heappush(remaining_sushi[t],-d)

eat_candidate_sushi = []
for i in remaining_sushi.keys():
    if not i in sushi_types:
        heapq.heappush(eat_candidate_sushi,remaining_sushi[i][0])

sushi_kinds=len(sushi_types)
val=v
ans=v+sushi_kinds**2

for i in range(sushi_kinds+1,k+1):
    if sushies and eat_candidate_sushi:
        v-=heapq.heappop(sushies)
        v-=heapq.heappop(eat_candidate_sushi)
        ans=max(ans,v+i**2)
    else:
        break

print(ans)

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

N, K = map(int, input().split())

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

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
for _ in range(K-used):
    d, t = heappop(not_used)
    su -= d
    heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    heappush(now_used, (-a, grp))
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

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

N, K = map(int, input().split())

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

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
for _ in range(K-used):
    d, t = heappop(not_used)
    su -= d
    heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

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

N, K = map(int, input().split())

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

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
if K-used:
    for _ in range(K-used):
        d, t = heappop(not_used)
        su -= d
        heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

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

n,k=map(int,input().split())
l=[list(map(int,input().split())) for _ in range(n)]
l.sort(key=lambda x:-x[1])

sushi_types = set()
sushies = []
v=0

for i in range(k):
    t,d=l[i]
    v+=d
    if t in sushi_types:
        heapq.heappush(sushies,d)
    else:
        sushi_types.add(t)

remaining_sushi=dd(list)
for i in range(k,n):
    t,d=l[i]
    if not t in sushi_types:
        heapq.heappush(remaining_sushi[t],-d)

l=[0]*(k+1)
current_types=len(sushi_types)
l[current_types]=v+current_types**2
ans=l[current_types]

for i in range(len(sushi_types)+1,k+1):
    t,add_target=min(remaining_sushi.items(),key=lambda x:x[1][0])
    sushi_types.add(t)
    remove_target=heapq.heappop(sushies)
    l[i]=l[i-1]-(i-1)**2-add_target[0]-remove_target+i**2
    ans=max(ans,l[i])
    del remaining_sushi[t]

print(ans)