output_description: Print the maximum satisfaction that you can obtain.
  
  * * *
input_description: Input is given from Standard Input in the following format:
  
      
      
      N K
      t_1 d_1
      t_2 d_2
      .
      .
      .
      t_N d_N
attempt: import sys
  
  fin = sys.stdin.readline
  from heapq import heappop, heappush
  from collections import defaultdict
  
  
  N, K = [int(elem) for elem in fin().split()]
  sushi_info = sorted(
      tuple(tuple(int(elem) for elem in fin().split()) for _ in range(N)),
      key=lambda x: -x[1],
  )
  
  cur_eaten = []
  cur_rest = []
  count = defaultdict(int)
  for t, d in sushi_info[:K]:
      heappush(cur_eaten, (d, t))
      count[t] += 1
  for t, d in sushi_info[K:]:
      heappush(cur_rest, (-d, t))
  
  num_distinctive_kinds = sum(v >= 1 for v in count.values())
  satisfactions = [sum(d for _, d in sushi_info[:K]) + num_distinctive_kinds**2]
  
  cur_num_kinds = num_distinctive_kinds
  for cur_num_kinds in range(num_distinctive_kinds + 1, K + 1):
      while cur_eaten and count[cur_eaten[0][1]] <= 1:
          heappop(cur_eaten)
      while cur_rest and count[cur_rest[0][1]] > 0:
          heappop(cur_rest)
      if (not cur_eaten) or (not cur_rest):
          break
      ditched_d, ditched_t = heappop(cur_eaten)
      added_d, added_t = heappop(cur_rest)
      added_d *= -1
      count[ditched_t] -= 1
      count[added_d] += 1
      # cur_num_kinds += 1
      satisfactions.append(
          satisfactions[-1] - ditched_d + added_d + cur_num_kinds * 2 - 1
      )
  
  # print(satisfactions)
  print(max(satisfactions))
  
problem_description: Statement
  
  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.
samples: [{"input": "5 3\n    1 9\n    1 7\n    2 6\n    2 5\n    3 1", "output": "26\n    \n\nIf you eat Sushi 1,2 and 3:\n\n  * The base total deliciousness is 9+7+6=22.\n  * The variety bonus is 2*2=4.\n\nThus, your satisfaction will be 26, which is optimal.\n\n* * *"}, {"input": "7 4\n    1 1\n    2 1\n    3 1\n    4 6\n    4 5\n    4 5\n    4 5", "output": "25\n    \n\nIt is optimal to eat Sushi 1,2,3 and 4.\n\n* * *"}, {"input": "6 5\n    5 1000000000\n    2 990000000\n    3 980000000\n    6 970000000\n    6 960000000\n    4 950000000", "output": "4900000016\n    \n\nNote that the output may not fit into a 32-bit integer type."}]