Problem p03148 (python)

Fixed submission:
# -*- coding: utf-8 -*-
from collections import deque
n,k = map(int, input().split())
td = []

for _ in range(n):
    t,d = map(int, input().split())
    td.append((t,d))

td.sort(reverse=True, key=lambda x:x[1])

used = set()
q = []
res = 0
kind = 0
for i in range(k):
    t,d = td[i]
    if t not in used:
        kind += 1
        used.add(t)
    else:
        q.append((t,d))
    res += d
res += kind**2
q = deque(sorted(q, key=lambda x:x[1]))

tmp = res
for i in range(k,n):
    if len(q)==0:
        break
    t2,d2 = td[i]
    if t2 in used:
        continue
    used.add(t2)
    t,d = q.popleft()

    r = tmp - d + d2
    r -= kind**2
    kind += 1
    r += kind**2
    res = max(res, r)
    tmp = r

print(res)

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
from collections import deque
n,k = map(int, input().split())
td = []

for _ in range(n):
    t,d = map(int, input().split())
    td.append((t,d))

td.sort(reverse=True, key=lambda x:x[1])

used = set()
q = []
res = 0
kind = 0
for i in range(k):
    t,d = td[i]
    if t not in used:
        kind += 1
        used.add(t)
    else:
        q.append((t,d))
    res += d
res += kind**2
q = deque(sorted(q, key=lambda x:x[1]))

tmp = res
for i in range(k,n):
    t2,d2 = td[i]
    if t2 in used:
        continue
    t,d = q.popleft()

    r = tmp - d + d2
    r -= kind**2
    kind += 1
    r += kind**2
    res = max(res, r)
    tmp = r
    if len(q)==0:
        break

print(res)

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
n,k = map(int, input().split())
o = {}

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

tl = list(o.keys())
tl.sort(reverse=True,key=lambda x:o[x])
# print(tl)
for t in tl:
    o[t].sort(reverse=True)
# print(o)
res = -1
sel = {}
for t in tl:
    if len(o[t])<=k:
        sel[t] = len(o[t])
        k -= sel[t]
    else:
        sel[t] = k
        k = 0
    if k==0:
        break

# print(sel)

def calc(sel):
    res = 0
    cnt = 0
    for t in sel:
        if sel[t]>0:
            cnt += 1
        for i in range(sel[t]):
            res += o[t][i]
    return res+cnt*cnt

idx = None
for i in range(len(tl)-1):
    if tl[i+1] not in sel:
        idx = i
        break

if idx is None:
    print(calc(sel))
    exit()

# print(idx)

res = calc(sel)
for i in range(idx+1,len(tl)):
    while sel[tl[idx]]==1:
        idx -= 1

    sel[tl[idx]] -= 1
    sel[i] = 1
    # res2 = res + o[i][0] - o[tl[idx]][-1]
    # res = max(res, res2)
    res = max(res, calc(sel))

print(res)

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
n,k = map(int, input().split())
o = {}

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

tl = list(o.keys())
tl.sort(reverse=True,key=lambda x:o[x])
# print(tl)
for t in tl:
    o[t].sort(reverse=True)
# print(o)
res = -1
sel = {}
for t in tl:
    if len(o[t])<=k:
        sel[t] = len(o[t])
        k -= sel[t]
    else:
        sel[t] = k
        k = 0
    if k==0:
        break

# print(sel)

def calc(sel):
    res = 0
    cnt = 0
    for t in sel:
        if sel[t]>0:
            cnt += 1
        for i in range(sel[t]):
            res += o[t][i]
    return res+cnt*cnt

idx = None
for i in range(len(tl)-1):
    if tl[i+1] not in sel:
        idx = i
        break

if idx is None:
    print(calc(sel))
    exit()

# print(idx)

res = calc(sel)
for i in range(idx+1,len(tl)):
    while sel[tl[idx]]==1:
        idx -= 1

    sel[tl[idx]] -= 1
    sel[i] = 1
    res2 = res + o[i][0] - o[tl[idx]][-1]
    res = max(res, res2)
    # res = max(res, calc(sel))

print(res)