Problem:


Score : 400 points

Problem 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.

Constraints

1 \leq K \leq N  \leq 10^5

1 \leq t_i \leq N

1 \leq d_i \leq 10^9

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

Output

Print the maximum satisfaction that you can obtain.

Sample Input 1

5 3
1 9
1 7
2 6
2 5
3 1

Sample Output 1

26

If you eat Sushi 1,2 and 3:

The base total deliciousness is 9+7+6=22.

The variety bonus is 2*2=4.

Thus, your satisfaction will be 26, which is optimal.

Sample Input 2

7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5

Sample Output 2

25

It is optimal to eat Sushi 1,2,3 and 4.

Sample Input 3

6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000

Sample Output 3

4900000016

Note that the output may not fit into a 32-bit integer type.



Sample input:
5 3
1 9
1 7
2 6
2 5
3 1


Sample output:
26

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    #  max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    #  content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx+1)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    #  for combination in combination_list:
    #      type_bonus = len(combination) * len(combination)

    #      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))
    #      if len(current_options) + len(max_values) < K:
    #          # そもそも不可能な組み合わせ
    #          continue
    #      elif len(current_options) + len(max_values) == K:
    #          results.append(sum(current_options) + sum(max_values) + type_bonus)
    #          continue
    #      current_options.sort(reverse=True)
    #      results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx+1)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    #  for combination in combination_list:
    #      type_bonus = len(combination) * len(combination)

    #      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))
    #      if len(current_options) + len(max_values) < K:
    #          # そもそも不可能な組み合わせ
    #          continue
    #      elif len(current_options) + len(max_values) == K:
    #          results.append(sum(current_options) + sum(max_values) + type_bonus)
    #          continue
    #      current_options.sort(reverse=True)
    #      results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx+1)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    for combination in combination_list:
        type_bonus = len(combination) * len(combination)

        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))
        if len(current_options) + len(max_values) < K:
            continue
        elif len(current_options) + len(max_values) == K:
            results.append(sum(current_options) + sum(max_values) + type_bonus)
            continue
        current_options.sort(reverse=True)
        results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx+1)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    for combination in combination_list:
        type_bonus = len(combination) * len(combination)

        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))
        if len(current_options) + len(max_values) < K:
            continue
        elif len(current_options) + len(max_values) == K:
            results.append(sum(current_options) + sum(max_values) + type_bonus)
        current_options.sort(reverse=True)
        results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx+1)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    for combination in combination_list:
        type_bonus = len(combination) * len(combination)

        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))
        if len(current_options) + len(max_values) < K:
            continue
        current_options.sort(reverse=True)
        results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import numpy as np


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())
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    neta_type_list = set([i[0] for i in TD])
    #  max_value_per_type = [[t, max([td[1] for td in TD if td[0] == t])]for t in neta_type_list]
    #  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 neta_type_list]
    max_value_per_type = [[i[0], max(i[1]), np.argmax(i[1])] for i in content_per_type]
    content_per_type = [[i[0], [j for idx, j in enumerate(i[1]) if idx != m[2]]] for i, m in zip(content_per_type, max_value_per_type)]

    combination_list = [list(itertools.combinations(neta_type_list, idx)) for idx, i in enumerate(neta_type_list)]
    combination_list = list(itertools.chain.from_iterable(combination_list))

    results = [0]
    for combination in combination_list:
        type_bonus = len(combination) * len(combination)

        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))
        if len(current_options) + len(max_values) < K:
            continue
        current_options.sort(reverse=True)
        results.append(sum(current_options[:K-len(combination)]) + sum(max_values) + type_bonus)






    #  for i in range(1, K+1):
    #      #  print('\n---------\n',i,'回目')
    #      # iを種類数とみなす
    #      # 何を選んだかは関係ない
        
    #      # 種類の全組み合わせを求める
    #      #  print(i)
    #      if len(neta_type_list) < i:
    #          break
    #      combinations = list(itertools.combinations(neta_type_list, 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]
    #          #  if len([i[1] for i in content_per_type 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()

Submission status: Wrong Answer

Solution code:
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:
            break 
        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(max(results))
 
 
main()

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  # sushis = sorted([tuple(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1])
  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = []
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types += [c[0]]
  types = set(types)

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for l in leave:
    if len(types) >= types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1
      types.add(l[0])

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])


if __name__=='__main__':
  main()

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = set()
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types = types.union({c[0]})

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for i, l in enumerate(leave):
    if i > types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))
  ps = base(curr) + len(types)**2
  di[len(types)] = ps
  for i in range(len(types)+1, K+1):
    if i <= types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      remove = curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      del leave[:j]
      curr += [leave.pop(0)]
    else:
      break
    di[str(i)] = (ps - remove[1] + curr[-1][1]) + i**2 - (i-1)**2
    ps = di[str(i)]
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i < types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      del leave[:j]
      curr += [leave.pop(0)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])