Problem: There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i. You are choosing K among these N pieces to eat. Your "satisfaction" here will be calculated as follows:    * The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".   * The base total deliciousness is the sum of the deliciousness of the pieces you eat.   * The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.  You want to have as much satisfaction as possible. Find this maximum satisfaction.

Submission attempts (in order):
# -*- coding: utf-8 -*-

def popMaxOishisa(sushis):
    t, d = max(sushis["max"].items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


def popMinOishisa_nonUniqueNeta(sel_sushis):
    t, d = min(sel_sushis["min"].items(), key=lambda i:i[1])
    sel_sushis[t].pop()
    if len(sel_sushis[t]) == 1:
        sel_sushis["min"].pop(t)
    else:
        sel_sushis["min"][t] = sel_sushis[t][-1]
    return t, d


def popMaxOishisa_uniqueNeta(sushis, netas):
    max_dict = {k:v for k,v in sushis["max"].items() if not k in netas}
    t, d = max(sushis["max"].items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


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


sushis = {}
sel_sushis = {}
for i in range(n):
    t, d = map(int, input().split())
    if not t in sushis:
        sushis[t] = [d]
    else:
        sushis[t].append(d)


for type in sushis:
    sushis[type].sort()
sushis["max"] = {t:sushis[t][-1] for t in sushis}

tMax = len(sushis) - 1


sum_oishisa = 0
for i in range(k):
    t, d = popMaxOishisa(sushis)
    if not t in sel_sushis:
        sel_sushis[t] = [d]
    else:
        sel_sushis[t].append(d)
    sum_oishisa += d
sel_sushis["min"] = {t:sel_sushis[t][-1] for t in sel_sushis if len(sel_sushis[t]) > 1}

tMin = len(sel_sushis) - 1
max_score = tMin**2 + sum_oishisa

for nNeta in range(tMin+1, tMax + 1, 1):

    # Find sushi having minimum oishisa and being not unique for neta
    t1, d1 = popMinOishisa_nonUniqueNeta(sel_sushis)
    sum_oishisa -= d1

    # Find sushi having maximum oishisa and unique for neta
    t2, d2 = popMaxOishisa_uniqueNeta(sushis, list(sel_sushis.keys()))
    sum_oishisa += d2

    if not t1 in sushis:
        sushis[t1] = [d1]
    else:
        sushis[t1].append(d1)

    if not t2 in sel_sushis:
        sel_sushis[t2] = [d2]
    else:
        sel_sushis[t2].append(d2)

    score = nNeta**2 + sum_oishisa
    if score > max_score:
        max_score = score

print(max_score)

N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
M = len(set([t for t, d in TD]))
 
TD.sort(key=lambda x: x[1], reverse=True)
p, S, V = 0, [], set()
for t, d in TD[:K]:
    p += d
    if t in V:
        S.append(d)
    else:
        V.add(t)
p += len(V) ** 2
ans = p
for t, d in TD[K:]:
    if len(S) == 0:
        break
    if t not in V:
        p += 2 * len(V) + 1 - S.pop() + d
        V.add(t)
        ans = max(ans, p)
print(ans)

Verdicts: Runtime Error

Accepted