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):
import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(set(t1[0:K-1])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2,key=lambda x: x[1], reverse=True)
t2,d2 = zip(*td2)
s = set()
ans2 = 0
counter = 0
for i in range(N):
    if(len(s) == K):
        break
    if(t2[i] not in s):
        ans2 +=d2[i]  
        s.add(t2[i])
print(ans2)
ans2 = ans2 + K ** 2
print(max(ans1,ans2))

import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(set(t1[0:K-1])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2,key=lambda x: x[1], reverse=True)
t2,d2 = zip(*td2)
s = set()
ans2 = 0
counter = 0
for i in range(N):
    if(len(s) == K):
        break
    if(t2[i] not in s):
        ans2 +=d2[i]  
        s.add(t2[i])
ans2 = ans2 + K ** 2
print(max(ans1,ans2))

import collections

index = 0

def maxsushi(sushi_list,eaten_sushi):
    global index
    max_val = [0,0]
    for i in range(index,len(sushi_list)):
        if(sushi_list[i][1] not in eaten_sushi):
            tmp = sushi_list[i]
            max_val = sushi_list[i]
            index = i
            break
    return max_val

N,K = map(int,input().split())
sushi = []
umasa = []
for _ in range(N):
    i,j = map(int,input().split())
    sushi.append(i)
    umasa.append(j)
l = list(zip(umasa,sushi))
l = sorted(l,reverse=True)
sushi_list = list(l)
eat_sushi = sushi_list[0:K]
eaten_sushi = set()
double_sushi = []
for i in range(K):
    if(sushi_list[i][1] not in eaten_sushi):
        eaten_sushi.add(sushi_list[i][1])
    else:
        double_sushi.append((sushi_list[i]))
manzoku = 0
sushi_list = sushi_list[K:N]
for i in range(len(eat_sushi)):
    manzoku += eat_sushi[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku 
while double_sushi:
    tmp1 = double_sushi.pop()
    tmp2 = maxsushi(sushi_list,eaten_sushi)
    if(tmp2[0] == 0 and tmp2[1] == 0):
        break
    eaten_sushi.add(tmp2[1])
    tmp_manzoku = tmp_manzoku + tmp2[0] - tmp1[0] + len(eaten_sushi) ** 2 - (len(eaten_sushi) - 1) ** 2
    if (tmp_manzoku > manzoku):
        manzoku = tmp_manzoku
print(manzoku)

Verdicts: Wrong Answer

Wrong Answer

Accepted