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):
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())
 

li.sort(reverse=True,key=lambda x:x[1])

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))



#-----------------------------------------
best_list=[]
not_best_list=li.copy()

for i in li:
    if i[0] not in [best_list[p][0] for p in range(len(best_list))]:
        best_list.append(i)
        
for i in best_list:
    not_best_list.remove(i)

#-----------------------------------------
s2=min(find_num_neta(list(range(n)),li),k)
s1=find_num_neta(list(range(k)),li)

oishisa=[0 for i in range(n)]
oishisa[s1]=sum([li[i][1] for i in range(k) ])
score=oishisa[s1]+s1**2

for i in range(s1+1,s2+1):
    oishisa[i]=oishisa[i-1]+best_list[i-1][1]-not_best_list[k-i][1]
    new_score=oishisa[i]+i**2
    if score<new_score:
        score=new_score
     
    
print(new_score)

n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())

li.sort(reverse=True,key=lambda x:x[1])

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

#-----------------------------------------
best_list=[]
best_list_neta=[]
not_best_list=[]

for i in li:
    if i[0] in best_list_neta:
        not_best_list.append(i[1])
    else:
        best_list.append(i[1])
        best_list_neta.append(i[0])


s2=min(len(best_list_neta),k)
s1=find_num_neta(list(range(k)),li)


score=-1
oishisa=[0 for i in range(n)]
oishisa[s1]=sum([li[p][1] for p in range(k)])

for i in range(s1+1,s2+1):
    oishisa[i]=oishisa[i-1]-not_best_list[k-i]+best_list[i-1]
    
    new_score=i**2+oishisa[i]

    if score<new_score:
        score=new_score
        
    
print(score)

# -*- coding: utf-8 -*-
"""
Created on Fri Feb 22 08:12:39 2019

@author: matsuda
"""

n,k=3,0
li=[[1,9],[3,9],[2,6]]

"""
n,k=7,4
li=[[1,1],
[2,1],
[3,1],
[4,6],
[4,5],
[4,5],
[4,5]]

"""

n,k=6,5
li=[[5,1000000000],
[2,990000000],
[3,980000000],
[6,970000000],
[6,960000000],
[4,950000000]]

#####################
li.sort(reverse=True,key=lambda x:x[1])

def calc_score(li,choice):
    score=0
    k=[]
    for i in choice:
        score+=li[i][1]
        k.append(li[i][0])        
    score+=len(set(k))**2
    return score

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

first_choice=list(range(k))
score=calc_score(li,first_choice)

li_choice_set=set([li[i][0] for i in first_choice])



for i in range(find_num_neta(first_choice,li)+1,min(find_num_neta(list(range(n)),li),k)+1):
  
    new_choice=first_choice.copy()
    change=i-find_num_neta(first_choice,li)
    
    can_append=change
    must_delete=change
    

    
    
    for j in range(k,n):
        if li[j][0] not in li_choice_set:
            new_choice.append(j)
            can_append+=-1
            if can_append==0:
                break
            
            
    for j in range(k-1,-1,-1):

        if li[j][0] in [li[p][0] for p in range(j)]:
            new_choice.remove(j)
            must_delete+=-1
            if must_delete==0:
                break
        
    
    new_score=calc_score(li,new_choice)

    
    if score<new_score:
        score=new_score
        
    
print(score)

n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())



def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))
#-----------------------------------------

li.sort(reverse=True,key=lambda x:x[1])
s2=min(find_num_neta(list(range(n)),li),k)
s1=find_num_neta(list(range(k)),li)
oishisa=sum([li[i][1] for i in range(k)])
score=oishisa+s1**2

#-----------------------------------------
li.sort(key=lambda x:x[0])
best_list=[]
not_best_list=[]


m=-1
for i in li:
    if m==i[0]:
        not_best_list.append(i)
    else:
        best_list.append(i)
        m=i[0]
        
not_best_list.sort(reverse=True,key=lambda x:x[1])
best_list.sort(reverse=True,key=lambda x:x[1])



 
 
for i in range(s1+1,s2+1):
    oishisa=oishisa+best_list[i-1][1]-not_best_list[k-i][1]
    
    new_score=i**2+oishisa
    
    if score<new_score:
        score=new_score
        
    
print(score)

Verdicts: Runtime Error

Runtime Error

Wrong Answer

Accepted