Buggy python submission (verdict: Time Limit Exceeded):
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    TD = []
    for i in range(N):
        TD.append(LI())
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in set([i[0] for i in TD])]
    #  print('max_value_per_type', max_value_per_type)
    content_per_type = [[t, [i[1] for i in TD if i[0] == t]] for t in set(i[0] for i in TD)]
    #  print(content_per_type)
    max_idx_per_type = [i[1].index([mvpt[1] for mvpt in max_value_per_type if mvpt[0] == i[0]][0]) for i in content_per_type]
    #  print('max_idx_per_type', max_idx_per_type)
    for idx, i in enumerate(max_idx_per_type):
        del content_per_type[idx][1][i]
    #  remaining_values_per_type = [[c[0], [i[0] for i in content_per_type if ]] for c in content_per_type]
    #  print('各タイプから最大値を抜きました', content_per_type)
    results = [0]
    for i in range(1, K+1):
        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        if len(set([i[0] for i in TD])) < i:
            continue
        combinations = list(itertools.combinations(set([i[0] for i in TD]), i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            max_values = [i[1] for i in max_value_per_type if i[0] in combination]
            current_options = [i[1] for i in content_per_type if i[0] in combination]
            current_options = list(itertools.chain.from_iterable(current_options))
            current_options.sort(reverse=True)
            #  print('最大値だけとった', base_point)
            #  print('残りの選択肢', current_options)
            results.append(sum(current_options[:K-i]) + sum(max_values) + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()


Fixed submission:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()

Buggy python submission (verdict: Time Limit Exceeded):
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    TD = []
    for i in range(N):
        TD.append(LI())
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    results = []
    max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in set([i[0] for i in TD])]
    #  print('max_value_per_type', max_value_per_type)
    content_per_type = [[t, [i[1] for i in TD if i[0] == t]] for t in set(i[0] for i in TD)]
    #  print(content_per_type)
    max_idx_per_type = [i[1].index([mvpt[1] for mvpt in max_value_per_type if mvpt[0] == i[0]][0]) for i in content_per_type]
    #  print('max_idx_per_type', max_idx_per_type)
    for idx, i in enumerate(max_idx_per_type):
        del content_per_type[idx][1][i]
    #  remaining_values_per_type = [[c[0], [i[0] for i in content_per_type if ]] for c in content_per_type]
    #  print('各タイプから最大値を抜きました', content_per_type)
    for i in range(1, K+1):
        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(set([i[0] for i in TD]), i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            max_values = [i[1] for i in max_value_per_type if i[0] in combination]
            current_options = [i[1] for i in content_per_type if i[0] in combination]
            current_options = list(itertools.chain.from_iterable(current_options))
            current_options.sort(reverse=True)
            #  print('最大値だけとった', base_point)
            #  print('残りの選択肢', current_options)
            results.append(sum(current_options[:K-i]) + sum(max_values) + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()


Fixed submission:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()

Buggy python submission (verdict: Time Limit Exceeded):
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    TD = []
    for i in range(N):
        TD.append(LI())
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    results = []
    max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in set([i[0] for i in TD])]
    #  print('max_value_per_type', max_value_per_type)
    content_per_type = [[t, [i[1] for i in TD if i[0] == t]] for t in set(i[0] for i in TD)]
    #  print(content_per_type)
    max_idx_per_type = [i[1].index([mvpt[1] for mvpt in max_value_per_type if mvpt[0] == i[0]][0]) for i in content_per_type] 
    #  print('max_idx_per_type', max_idx_per_type)
    for idx, i in enumerate(max_idx_per_type):
        del content_per_type[idx][1][i]
    #  remaining_values_per_type = [[c[0], [i[0] for i in content_per_type if ]] for c in content_per_type]
    #  print('各タイプから最大値を抜きました', content_per_type)
    for i in range(1, K+1):
        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 最小種類数のときの最大値、最大種類数のときの最大値を計算して比較する
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(set([i[0] for i in TD]), i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            base_point = 0
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            required_neta_count = K
            max_values = [i for i in max_value_per_type if i[0] in combination]
            base_point += sum([i[1] for i in max_values])
            current_options = [i[1] for i in content_per_type if i[0] in combination]
            current_options = list(itertools.chain.from_iterable(current_options))
            current_options.sort(reverse=True)
            #  print('最大値だけとった', base_point)
            #  print('残りの選択肢', current_options)
            #  #  current_options = [i for i in current_options if not(i[0] in combination and i[1] == [i[0] for i in ])]
            #  # 各タイプの最大値をまず選び、current_optionsから排除
            #  for t in combination:
            #      max_value = [i for i in max_value_per_type if i[0] == t][0][1]
            #      base_point += max_value

            #      # max_valueは選択肢から取り除いておく
            #      current_options = [i for i in current_options if not (i[0] == t and i[1] == max_value)]
            #      #  print('選んだ寿司を選択肢から取り除いた', current_options)

            #  print('あと何個ネタを取るべきか? ', required_neta_count, '個')
            # あとは残りの選択肢から大きいものを選ぶ
            #  print('current_options, requried_neta_count', current_options, required_neta_count)
            results.append(sum(current_options[:K-i]) + base_point + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()


Fixed submission:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()

Buggy python submission (verdict: Runtime Error):
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    TD = []
    for i in range(N):
        TD.append(LI())
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    results = []
    for i in range(1, K+1):
        if i == K:
            # 全タイプの最大値を取る
            base_point = 0
            for t in [set(i[0] for i in TD)]:
                base_point += max([i for i in TD if i[0] == t])

            results.append(base_point + i * i)


        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 最小種類数のときの最大値、最大種類数のときの最大値を計算して比較する
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(set([i[0] for i in TD]), i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            base_point = 0
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            required_neta_count = K
            # 各タイプの最大値をまず選び、current_optionsから排除
            for t in combination:
                max_value = max([option[1] for option in current_options if option[0] == t])
                base_point += max_value

                # max_valueは選択肢から取り除いておく
                current_options = [i for i in current_options if not (i[0] == t and i[1] == max_value)]
                #  print('選んだ寿司を選択肢から取り除いた', current_options)

            #  print('あと何個ネタを取るべきか? ', required_neta_count, '個')
            # あとは残りの選択肢から大きいものを選ぶ
            #  print('current_options, requried_neta_count', current_options, required_neta_count)
            results.append(sum([i[1] for i in current_options[:K-i]]) + base_point + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()


Fixed submission:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()

Buggy python submission (verdict: Runtime Error):
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    for i in range(N):
        t, d = LI()
        T.append(t)
        D.append(d)
    #  print(T, '\n', D)
    TD = [[t, d] for t, d in zip(T, D)]
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    all_types = set(T)
    results = []
    for i in range(1, K+1):
        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 最小種類数のときの最大値、最大種類数のときの最大値を計算して比較する
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(all_types, i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            base_point = 0
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            required_neta_count = K
            # 各タイプの最大値をまず選び、current_optionsから排除
            for t in combination:
                base_point += max([option[1] for option in current_options if option[0] == t])

                # max_valueは選択肢から取り除いておく
                current_options = [i for i in current_options if not (i[0] == t and i[1] == max_value)]
                #  print('選んだ寿司を選択肢から取り除いた', current_options)

            #  print('あと何個ネタを取るべきか? ', required_neta_count, '個')
            # あとは残りの選択肢から大きいものを選ぶ
            #  print('current_options, requried_neta_count', current_options, required_neta_count)
            results.append(sum([i[1] for i in current_options[:K-i]]) + base_point + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()


Fixed submission:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()