Problem p03148 (python)

Fixed submission:
# 解説放送

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

e = [None for _ in range(N + 1)] # 種別の美味しさ
t_kinds_count = 0 # 何種類あるか
rest_length = 0 # 各種の2番手以降の個数の総数

tops = [] # 各種の最大美味しさ
c_tops = [0] # 各種の最大美味しさの累積和
rests = [] # 各種の2番手以降
c_rests = [0] # 各種の2番手以降の累積和

for _ in range(N):
    t, d = map(int, input().split())
    if e[t] is None:
        e[t] = []
        t_kinds_count += 1
    e[t].append(d)


for t in range(1, N + 1):
    if e[t] is None:
        continue
    e[t].sort(reverse=True)
    tops.append(e[t][0])
    rests.extend(e[t][1:])

tops.sort(reverse=True)
rests.sort(reverse=True)

s = 0
for tt in tops:
    s += tt
    c_tops.append(s)

s = 0
for rr in rests:
    s += rr
    c_rests.append(s)
    rest_length += 1

ans = 0
for yy in range(1, min(t_kinds_count, K) + 1):
    if 0 <= K - yy <= rest_length:
        ans = max(ans, c_tops[yy] + c_rests[K - yy] + yy * yy)

# print(e)
# print(tops)
# print(c_tops)
# print(rests)
# print(c_rests)
print(ans)

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

inf = float('inf')

n, k = map(int, input().split())

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

for key in dd:
    dd[key].sort(reverse=True)

h1 = []
h2 = []
for key in dd:
    for i, d in enumerate(dd[key]):
        if i == 0:
            h1.append(d)
        else:
            h2.append(d)

if h1:
    h1.sort(reverse=True)
if h2:
    h2.sort(reverse=True)

c1 = [0]
for hh in h1:
    c1.append(c1[-1] + hh)
if len(c1) < k:
    c1.extend([-inf] * (k - len(c1) + 1))

c2 = [0]
for hh in h2:
    c2.append(c2[-1] + hh)
if len(c2) < k:
    c2.extend([-inf] * (k - len(c2) + 1))

# print(h1, c1)
# print(h2, c2)

ans = -inf
for kind in range(1, k + 1):
    # print(kind, c1[kind], c2[k - kind])
    ans = max(ans, c1[kind] + c2[k - kind] + kind * kind)
print(ans)

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

inf = float('inf')

n, k = map(int, input().split())

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

for key in dd:
    dd[key].sort(reverse=True)

h1 = []
h2 = []
for key in dd:
    for i, d in enumerate(dd[key]):
        if i == 0:
            h1.append(d)
        else:
            h2.append(d)

h1.sort(reverse=True)
h2.sort(reverse=True)

c1 = [0]
for hh in h1:
    c1.append(c1[-1] + hh)
if len(c1) < k:
    c1.extend([-inf] * (k - len(c1) + 1))

c2 = [0]
for hh in h2:
    c2.append(c2[-1] + hh)
if len(c2) < k:
    c2.extend([-inf] * (k - len(c2) + 1))

# print(h1, c1)
# print(h2, c2)

ans = -inf
for kind in range(1, k + 1):
    # print(kind, c1[kind], c2[k - kind])
    ans = max(ans, c1[kind] + c2[k - kind] + kind * kind)
print(ans)

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

input = sys.stdin.readline

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

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, key=itemgetter(0), reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)