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 = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for _ in range(n)]


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



eated_kind = []
eated_kind_num = 0
eated = []
for t,d in td:
  if t not in eated_kind:
    eated_kind.append(t)
    eated_kind_num += 1
    eated.append([t,d])
    
    
    
    if eated_kind_num == k:
      td = td[:i+1]
      break

not_eated = [i for i in td if i not in eated]


if eated_kind_num < k:
  eated += not_eated[:k-eated_kind_num]
  not_eated = not_eated[k-eated_kind_num:]
  #print(eated)
  eated.sort(key = lambda x: x[1],reverse=True)
  
t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point

for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  eated = eated[:-1]
  point = max(point,point[-1]+t_point_sub+d_point_sub)
  eated_kind_num -= 1    
#print(eated)

print(point)

n,k = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for _ in range(n)]


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



eated_kind = []
eated_kind_num = 0
eated = []
for t,d in td:
  if t not in eated_kind:
    eated_kind.append(t)
    eated_kind_num += 1
    eated.append([t,d])
    
    
    
    if eated_kind_num == k:
      td = td[:i+1]
      break

not_eated = [i for i in td if i not in eated]


if eated_kind_num < k:
  eated += not_eated[:k-eated_kind_num]
  not_eated = not_eated[k-eated_kind_num:]
  #print(eated)
  eated.sort(key = lambda x: x[1],reverse=True)
  
t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point

for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  eated = eated[:-1]
  point = max(point,point+t_point_sub+d_point_sub)
  eated_kind_num -= 1
  if eated_kind_num == 1:
    break
#print(eated)

print(point)

n,k = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for _ in range(n)]


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


fix_eated_d=[]

change_eated_d=[]

eated_t = set()
d_point = 0


for t,d in td[:k]:
  
  
  if t in eated_t:
    change_eated_d.append(d)
    
  
  else:
    fix_eated_d.append(d)
    eated_t.add(t)
    
  d_point += d
  
point = len(eated_t)**2 + d_point


for t,d in td[k:]:
  
  if not change_eated_d:
    break
  
  
  if t in eated_t:
    continue
    
  
  eated_t.add(t)
  d_point += d - change_eated_d.pop()
  point = max(point, len(eated_t)**2 + d_point)
  
print(point)




Verdicts: Runtime Error

Runtime Error

Accepted