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
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, len(set([i[0] for i in TD])) + 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()

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
      leave[:j]
      curr += [leave.pop(0)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

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)

    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, len(set([i[0] for i in TD])) + 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

            required_neta_count = K
            max_values = [i 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([i[1] for i in max_values]) + type_bonus)

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


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])[::-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
      curr += [leave.pop(j)]
    else:
      break
    print(curr)
  print(di)
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

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)

    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])] 
    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
            # 各タイプの最大値をまず選び、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([i[1] for i in current_options[:K-i]]) + base_point + 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)

    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[1] for i in TD if i[0] == t])

            results.append(base_point + i * i)
            continue


        #  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()

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: # len(curr):
      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
      curr += [leave.pop(j)]
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

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 = LI()
        TD.append(td)
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    results = []
    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
            # 各タイプの最大値をまず選び、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()

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 = []
    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:
                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()

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 = []
    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)
    point = 0
    types = []
    estimated_values = []
    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:
                #  current_options = [i for i in TD if i[0] in combination]
                max_value = max([option[1] for option in current_options if option[0] == t])
                #  print('max_value', max_value)
                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)

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

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


main()

Submission status: Wrong Answer

Solution code:
import sys

def main():
  input = sys.stdin.readline
  n, k = map(int, input().split())
  sushi = []
  for _ in range(n):
    t, d = map(int, input().split())
    sushi.append([t,d])
  
  sushi = sorted(sushi, key=lambda x:-x[1])
  point = 0
  d = []
  t = []
  for s in sushi[:k]:
    if s[0] in t:
      d.append(s[1])
    else:
      point += s[1]
      t.append(s[0])
  d = sorted(d)
  t = set(t)
  rest = [s[1] for s in sushi[k:] if s[0] not in t]
  ans = point + sum(d) + len(t)**2
  m = min(len(rest), len(d))
  m = max(0, m)
  for i in range(m):
    point += rest[i]
    res = point + sum(d[i+1:]) + (len(t) + i + 1)**2
    ans = max(res, ans)
  print(ans)


if __name__ == '__main__':
  main()

Submission status: Wrong Answer

Solution code:
#ABC116D-VariousSushi.py

INPUT1 = [int(n) for n in input().split()]

N, K = INPUT1[0], INPUT1[1]
Parameters = []
Score, Test = 0, 0
flag1, flag2 = 0, 0


for i in range(N):
    INPUT2 = [int(m) for m in input().split()]
    Parameters.append((INPUT2[0], INPUT2[1]))

Parameters = sorted(Parameters, key=lambda x: x[1], reverse=True)

def SumScore(List, K):
    taste = 0
    kind = []
    for i in range(K):
        taste += List[i][1]
        kind.append(List[i][0])
    return taste + len(set(kind))**2

def SearchNew(List, K):
    kind = []
    newkind = -1
    for i in range(K):
        kind.append(List[i][0])
    kindset = set(kind)
    for i in range(len(List)-K):
        if List[K+i][0] not in kindset:
            newkind = K+i
            break
    return newkind

def SearchDouble(List, K):
    result = -1
    for i in range(K):
        leftkind = []
        for j in range(K):
            leftkind.append(List[j][0])
        testkind = List[K-1-i][0]
        leftkind.pop(K-1-i)
        if testkind in leftkind:
            result = K-1-i
            break
    return result

Score = SumScore(Parameters, K)

for t in Parameters:
    i = SearchNew(Parameters, K)
    j = SearchDouble(Parameters, K)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    Test = SumScore(Parameters, K)
    if Score <= Test:
        Score = Test
    print(i, j, t)

print(Score)

Submission status: Wrong Answer

Solution code:
a = list(map(lambda x: int(x), input().split()))
N = a[0]
K = a[1]
sushis = []
for i in range(1, N):
  sushis.append(list(map(lambda x: int(x), input().split())))

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

netas = set()
duplicates = []

choices = sushis[:K]

for choice in choices:
  if choice[0] in netas:
    duplicates.append(choice[1])
  else:
    netas.add(choice[0])

p = sum(map(lambda x: x[1], choices)) + len(netas)**2
answer = p

for sushi in sushis[K:]:
  if sushi[0] not in netas and duplicates:
    netas.add(sushi[0])
    duplicate = duplicates.pop()
    p += (sushi[1] - duplicate) + 2*len(netas) - 1
    answer = max(answer, p)

print(answer)

Submission status: Wrong Answer

Solution code:
def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]
 
    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None
 
N, K = map(int, input().split())
menu_ = []
 
for i in range(N):
    t, d = map(int, input().split())
    menu_.append((t, d))
 
menu = sorted(menu_, key = lambda x:x[1], reverse=True)
p = sum([d for (t, d) in menu[:K]]) + len(set([t for (t, d) in menu[:K]]))**2
 
selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]
ans = p
for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
        p += (m[1] - d[1]) + 2 * len(set(t_list)) - 1 
        #print(ans, p)
        ans = max(ans, p)
 
print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta and duplicates:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) - 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans)

Submission status: Wrong Answer

Solution code:
def manzoku(s):
    return sum([d for (t, d) in s]) + len(set([t for (t, d) in s]))**2

def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]

    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
ans = manzoku(menu[:K])

selected_menu = menu[:K]
#t_list = [t for (t, d) in selected_menu]
#d_list = [d for (t, d) in selected_menu]

for m in menu[K:]:
    t_list = [t for (t, d) in selected_menu]
    d_list = [d for (t, d) in selected_menu]
    
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
    p = manzoku(selected_menu)
    #print(ans, p)
    ans = max(ans, p)

print(ans)

Submission status: Wrong Answer

Solution code:
def manzoku(s):
    return sum([d for (t, d) in s]) + len(set([t for (t, d) in s]))**2

def seek_dup(s):
    t_list = [t for (t, d) in s]
    #d_list = [d for (t, d) in s]

    for t, d in reversed(s):
        if t_list.count(t) > 1:
            return (t, d)
        else:
            continue
        
    return None

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))

sorted(menu, key = lambda x:x[1], reverse=True)
ans = manzoku(menu[:K])

selected_menu = menu[:K]
t_list = [t for (t, d) in selected_menu]
d_list = [d for (t, d) in selected_menu]


for m in menu[K:]:
    if m[0] not in t_list:
        d = seek_dup(selected_menu)
        if d == None:
            continue
        
        selected_menu.remove(d)
        selected_menu.append(m)
    
    p = manzoku(selected_menu)
    #print(ans, p)
    ans = max(ans, p)

print(ans)