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):
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, vari


selected = sushis[:K]
unSelected = sushis[K:]

p, vari = calcPoint(selected)
cand = [p]

va = len(vari)
v = [0]*(N+1)
for i, _ in selected:
    v[i] += 1

can_del = []
for i, v in reversed(selected):
    if vari[i] > 2:
        vari[i] -= 1
        can_del.append(v)

can_in = []
for i, v in unSelected:
    if not i in vari:
        vari.add(i)
        can_in.append(v)

for od, nd in zip(can_del, can_in):
    nextp = p-od+nd-(vari)**2+(vari+1)**2
    cand.append(nextp)
    vari += 1
print(max(cand))


# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, len(vari)


selected = sushis[:K]
rest_shushi = sushis[K+1:]
p, nv = calcPoint(selected)

vari = [0]*(N+1)
for i, v in selected:
    vari[i] += 1

for i, v in rest_shushi:
    if vari[i] >= 1:
        continue
    else:
        d = []
        for a, b in reversed(selected):
            if vari[a] > 1:
                next_p = p-b+v + (nv+1)**2 - nv**2
                if next_p > p or True:
                    vari[a] -= 1
                    vari[i] += 1
                    nv += 1
                    p = next_p
print(p)


# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, vari


selected = sushis[:K]
unSelected = sushis[K:]

p, vari = calcPoint(selected)
cand = [p]

va = len(vari)
v = [0]*(N+1)
for i, _ in selected:
    v[i] += 1

can_del = []
for i, d in reversed(selected):
    if v[i] >= 2:
        v[i] -= 1
        can_del.append(d)

can_in = []
for i, d in unSelected:
    if not i in vari:
        vari.add(i)
        can_in.append(d)

#print("del->", can_del, can_in)
for od, nd in zip(can_del, can_in):
    nextp = p-od+nd-(va)**2+(va+1)**2
    p = nextp
    cand.append(p)
    va += 1
print(max(cand))


Verdicts: Runtime Error

Wrong Answer

Accepted