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):
def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]
 
    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None
 
N, K = map(int, input().split())
menu_ = []
 
for i in range(N):
    t, d = map(int, input().split())
    menu_.append((t, d))
 
menu = sorted(menu_, key = lambda x:x[1], reverse=True)
p = sum([d for (t, d) in menu[:K]]) + len(set([t for (t, d) in menu[:K]]))**2
 
selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]
ans = p
for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
    p += (m[1] - d[1]) + 2 * len(set(t_list)) - 1 
    #print(ans, p)
    ans = max(ans, p)
 
print(ans) 

def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]

    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
ans = sum([d for (t, d) in menu[:K]]) + len(set([t for (t, d) in menu[:K]]))**2

selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]
p = ans
for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
    p += (m[1] - d[1]) + 2 * len(set(t_list)) - 1 
    #print(ans, p)
    ans = max(ans, p)

print(ans) 

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) + 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans) 


def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]

    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
ans = sum([d for (t, d) in menu[:K]]) + len(set([t for (t, d) in menu[:K]]))**2

selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]
p = ans
for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        #selected_menu.remove(d)
        #selected_menu.append(m)
    
    p += (m[1] - d[1]) + 2 * len(set(t_list)) + 1 
    #print(ans, p)
    ans = max(ans, p)

print(ans)    
    

import itertools

N, K = map(int, input().split())
menu = {}

t0 = time.time()
for i in range(N):
    t, d = map(int, input().split())
    menu[i] = (t, d)

def manzoku(s):
    return sum([menu[i][1] for i in s]) + len(set([menu[i][0] for i in s]))**2

shushi = itertools.combinations(menu, K)
print(max([manzoku(c) for c in shushi]))



N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

menu = sorted(menu, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) + 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans)  

def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]
 
    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None
 
N, K = map(int, input().split())
menu_ = []
 
for i in range(N):
    t, d = map(int, input().split())
    menu_.append((t, d))
 
menu = sorted(menu_, key = lambda x:x[1], reverse=True)
p = sum([d for (t, d) in menu[:K]]) + len(set([t for (t, d) in menu[:K]]))**2
 
selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]
ans = p
for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list and len(set(t_list)) == len(t_list):
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
    p += (m[1] - d[1]) + 2 * len(set(t_list)) - 1 
    #print(ans, p)
    ans = max(ans, p)
 
print(ans) 

N, K = map(int, input().split())
menu_ = []

for i in range(N):
    t, d = map(int, input().split())
    menu_.append((t, d))

menu = sorted(menu_, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta and duplicates:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) - 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans)    
    

Verdicts: Runtime Error

Runtime Error

Runtime Error

Runtime Error

Runtime Error

Runtime Error

Runtime Error

Accepted