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 heapq
  from collections import defaultdict
  
  N, K = map(int, input().split())
  
  sushi = []
  for _ in range(N):
      kind, value = map(int, input().split())
      sushi.append((value, kind))
  sushi.sort()
  
  
  choice = []
  heapq.heapify(choice)
  
  current_sum = 0
  kinds = set()
  count_nitems_bykind = defaultdict(int)
  
  # pick K items with highest value
  for _ in range(K):
      item = sushi.pop()
      count_nitems_bykind[item[1]] += 1
      kinds.add(item[1])
      current_sum += item[0]
      heapq.heappush(choice, item)
  
  # initialize a dict to store max value that can be achieved by selecting x kinds of sushi
  # key: x, value: max value
  max_value = {len(kinds): current_sum}
  
  
  def replace():
      """
      Returns a next sushi that is new in kind and has maximum value
      """
  
      while sushi:
          candidate = sushi.pop()
          if candidate[1] not in kinds:
              kinds.add(candidate[1])
              return candidate
  
      return None
  
  
  while sushi and choice:
      # pick a sushi that has min value
      sushi_min = heapq.heappop(choice)
  
      # check if that sushi is a duplicate then if yes replace with a new sushi whose kind is new and value is max among such
      if count_nitems_bykind[sushi_min[1]] > 1:
  
          # pick a candidate to replace
          sushi_new = replace()
          if sushi_new is None:
              break
  
          count_nitems_bykind[sushi_min[1]] -= 1
          heapq.heappush(choice, sushi_new)
  
          # add a new result to dict
          current_sum = current_sum - sushi_min[0] + sushi_new[0]
          max_value[len(kinds)] = current_sum
  
      if len(kinds) == K:  # already iterated enough to find answer
          break
  
  ans = -1
  
  for key in max_value:
      ans = max(ans, max_value[key] + key**2)
  
  print(ans)
  
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."}]