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: Time Limit Exceeded

Solution code:
#ABC116D-VariousSushi.py

import sys
import numpy as np
from operator import itemgetter

input = sys.stdin.readline

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

N, K = INPUT1[0], INPUT1[1]
Parameters = [[None] * N for _ in range(N)]
Score= 0
kind = [None] * K
kindnumber = 0

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

Parameters = sorted(Parameters, key=itemgetter(1), reverse=True)

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

def SearchDouble(List, kind, J):
    result = -1
    testkind = kind[:]
    oldkindnumber = len(set(testkind))
    testkindPop = testkind.pop
    for i in range(J):
        testkindPop(J-1-i)
        newkindnumber = len(set(testkind))
        if oldkindnumber == newkindnumber:
            result = J-1-i
            break
        oldkindnumber = newkindnumber
    return result

for i in range(K):
    kind[i] = Parameters[i][0]
kindnumber = len(set(kind))
Score = np.sum(Parameters[:K], axis=0)[1] + kindnumber**2

Test1, Test2 = 0, Score

I, J = K, K

for t in range(N):
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, kind, I)
    j = SearchDouble(Parameters, kind, J)
    if i != -1 and j != -1:
        kindnumber += 1
        I, J = i, j
        kind[j] = Parameters[i][0]
        Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    else:
        break
    Test1 = Test2 - Parameters[i][1] + Parameters[j][1] + kindnumber * 2 - 1
    Test2 = Test1
    if Score <= Test2:
        Score = Test2

print(Score)

Submission status: Time Limit Exceeded

Solution code:
#ABC116D-VariousSushi.py

import numpy as np

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

N, K = INPUT1[0], INPUT1[1]
Parameters = [[None for col in range(2)] for row in range(N)]
Score= 0
kind = []
kindnumber = 0

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

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

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

def SearchDouble(List, kind, J):
    result = -1
    testkind = kind[:]
    oldkindnumber = len(set(testkind))
    testkindPop = testkind.pop
    for i in range(J):
        testkindPop(J-1-i)
        newkindnumber = len(set(testkind))
        if oldkindnumber == newkindnumber:
            result = J-1-i
            break
        oldkindnumber = newkindnumber
    return result

kindAppend = kind.append

for i in range(K):
    kindAppend(Parameters[i][0])
kindnumber = len(set(kind))
Score = np.sum(Parameters[:K], axis=0)[1] + kindnumber**2

Test1, Test2 = 0, Score

I, J = K, K

for t in Parameters:
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, kind, I)
    j = SearchDouble(Parameters, kind, J)
    if i != -1 and j != -1:
        kindnumber += 1
        I, J = i, j
        kind[j] = Parameters[i][0]
        Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    else:
        break
    Test1 = Test2 - Parameters[i][1] + Parameters[j][1] + kindnumber * 2 - 1
    Test2 = Test1
    if Score <= Test2:
        Score = Test2

print(Score)

Submission status: Time Limit Exceeded

Solution code:
import itertools
def gen_comb(n, r):
    return list(itertools.combinations([i for i in range(n)], r))

def calc_points(umapt, n_kinds):
    return sum(umapt)+n_kinds**2
  
# 貪欲法
def greed(name, umai, n, r):
    combs = gen_comb(n, r)
    max_val = max([calc_points([umai[i] for i in idx], len(set([name[i] for i in idx]))) for idx in combs])
    return max_val

N,K = [int(i) for i in input().split()]
name = [0]*N
umai = [0]*N
for i in range(N):
    name[i],umai[i] = [int(i) for i in input().split()]
print(greed(name, umai, N, K))

Submission status: Time Limit Exceeded

Solution code:
import itertools
def gen_comb(n, r):
    return list(itertools.combinations([i for i in range(n)], r))

def calc_points(umapt, n_kinds):
    return sum(umapt)+n_kinds**2
  
# 貪欲法
def greed(name, umai, n, r):
    combs = gen_comb(n, r)
    max_ = 0
    for idx in combs:
        umapt = [umai[i] for i in idx]
        n_kinds = len(set([name[i] for i in idx]))
        point = calc_points(umapt, n_kinds)
        if point > max_:
            max_ = point
    return max_

N,K = [int(i) for i in input().split()]
name = [0]*N
umai = [0]*N
for i in range(N):
    name[i],umai[i] = [int(i) for i in input().split()]
print(greed(name, umai, N, K))

Submission status: Time Limit Exceeded

Solution code:
#ABC116D-VariousSushi.py

import numpy as np
import bisect

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

N, K = INPUT1[0], INPUT1[1]
Parameters = [[None for col in range(2)] for row in range(N)]
Score= 0
kind = []
kindnumber = 0

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

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

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

def SearchDouble(List, K, kind):
    result = -1
    testkind = kind[:]
    oldkindnumber = len(set(testkind))
    testkindPop = testkind.pop
    for i in range(K):
        testkindPop(K-1-i)
        newkindnumber = len(set(testkind))
        if oldkindnumber == newkindnumber:
            result = K-1-i
            break
        oldkindnumber = newkindnumber
    return result

kindAppend = kind.append

for i in range(K):
    kindAppend(Parameters[i][0])
kindnumber = len(set(kind))
Score = np.sum(Parameters[:K], axis=0)[1] + kindnumber**2
kind.clear()

Test1, Test2 = 0, Score

for t in Parameters:
    for i in range(K):
        kindAppend(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    else:
        break
    Test1 = Test2 - Parameters[i][1] + Parameters[j][1] + kindnumber * 2 - 1
    Test2 = Test1
    if Score <= Test2:
        Score = Test2
    kind.clear()

print(Score)

Submission status: Time Limit Exceeded

Solution code:
#ABC116D-VariousSushi.py

import numpy as np

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

N, K = INPUT1[0], INPUT1[1]
Parameters = []
Score= 0
kind = []
kindnumber = 0

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

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

def binanrySearch(kind, target):
    testkind = kind[:]
    testkind.sort()
    left, right = 0, len(testkind)-1
    while(right >= left):
        mid = int((right + left) / 2)
        if testkind[mid] == target:
            return True
        elif testkind[mid] > target:
            right = mid - 1
        elif testkind[mid] < target:
            left = mid + 1
    return False

def SearchNew(List, K, kind):
    newkind = -1
    for i in range(len(List)-K):
        if binanrySearch(kind, List[K+i][0]) == False:
            newkind = K+i
            break
    return newkind

def SearchDouble(List, K, kind):
    result = -1
    testkind = kind[:]
    testkindPop = testkind.pop
    for i in range(K):
        testkindPop(K-1-i)
        if binanrySearch(testkind, List[K-1-i][0]) == True:
            result = K-1-i
            break
    return result

kindAppend = kind.append

for i in range(K):
    kindAppend(Parameters[i][0])
kindnumber = len(set(kind))
Score = np.sum(Parameters[:K], axis=0)[1] + kindnumber**2
kind.clear()

Test1, Test2 = 0, Score

for t in Parameters:
    for i in range(K):
        kindAppend(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    else:
        break
    Test1 = Test2 - Parameters[i][1] + Parameters[j][1] + kindnumber * 2 - 1
    Test2 = Test1
    if Score <= Test2:
        Score = Test2
    kind.clear()

print(Score)

Submission status: Time Limit Exceeded

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
kind = []
kindnumber = 0

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

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

def binanrySearch(kind, target):
    testkind = kind[:]
    testkind.sort()
    left, right = 0, len(testkind)-1
    while(right >= left):
        mid = int((right + left) / 2)
        if testkind[mid] == target:
            return True
        elif testkind[mid] > target:
            right = mid - 1
        elif testkind[mid] < target:
            left = mid + 1
    return False

def SumScore(List, K, kindnumber):
    taste = 0
    for i in range(K):
        taste += List[i][1]
    return taste + kindnumber * kindnumber

def SearchNew(List, K, kind):
    newkind = -1
    for i in range(len(List)-K):
        if binanrySearch(kind, List[K+i][0]) == False:
            newkind = K+i
            break
    return newkind

def SearchDouble(List, K, kind):
    result = -1
    testkind = kind[:]
    testkindPop = testkind.pop
    for i in range(K):
        testkindPop(K-1-i)
        if binanrySearch(testkind, List[K-1-i][0]) == True:
            result = K-1-i
            break
    return result

kindAppend = kind.append

for i in range(K):
    kindAppend(Parameters[i][0])
kindnumber = len(set(kind))
Score = SumScore(Parameters, K, kindnumber)
kind.clear()

for t in Parameters:
    for i in range(K):
        kindAppend(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    else:
        break
    Test = SumScore(Parameters, K, kindnumber)
    if Score <= Test:
        Score = Test
    kind.clear()

print(Score)

Submission status: Time Limit Exceeded

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
kind = []
kindnumber = 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 binanrySearch(kind, target):
    testkind = kind[:]
    testkind.sort()
    left, right = 0, len(testkind)-1
    while(right >= left):
        mid = int((right + left) / 2)
        if testkind[mid] == target:
            return True
        elif testkind[mid] > target:
            right = mid - 1
        elif testkind[mid] < target:
            left = mid + 1
    return False

def SumScore(List, K, kindnumber):
    taste = 0
    for i in range(K):
        taste += List[i][1]
    return taste + kindnumber * kindnumber

def SearchNew(List, K, kind):
    newkind = -1
    for i in range(len(List)-K):
        if binanrySearch(kind, List[K+i][0]) == False:
            newkind = K+i
            break
    return newkind

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

for i in range(K):
    kind.append(Parameters[i][0])
kindnumber = len(set(kind))
Score = SumScore(Parameters, K, kindnumber)
kind.clear()

for t in Parameters:
    for i in range(K):
        kind.append(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    else:
        break
    Test = SumScore(Parameters, K, kindnumber)
    if Score <= Test:
        Score = Test
    kind.clear()

print(Score)

Submission status: Time Limit Exceeded

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
kind = []
kindnumber = 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, kindnumber):
    taste = 0
    for i in range(K):
        taste += List[i][1]
    return taste + kindnumber * kindnumber

def SearchNew(List, K, kind):
    newkind = -1
    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, kind):
    result = -1
    for i in range(K):
        testkind = List[K-1-i][0]
        kind.pop(K-1-i)
        if testkind in kind:
            result = K-1-i
            break
    return result

for i in range(K):
    kind.append(Parameters[i][0])
kindnumber = len(set(kind))
Score = SumScore(Parameters, K, kindnumber)

for t in Parameters:
    for i in range(K):
        kind.append(Parameters[i][0])
    kindnumber = len(set(kind))
    i = SearchNew(Parameters, K, kind)
    j = SearchDouble(Parameters, K, kind)
    Parameters[i], Parameters[j] = Parameters[j], Parameters[i]
    if i != -1 and j != -1:
        kindnumber += 1
    Test = SumScore(Parameters, K, kindnumber)
    if Score <= Test:
        Score = Test
    kind.clear()

print(Score)

Submission status: Time Limit Exceeded

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('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: Time Limit Exceeded

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

Submission status: Time Limit Exceeded

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

Submission status: Time Limit Exceeded

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

import sys
input = sys.stdin.readline

if __name__=='__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)

  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 = list(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 += [l[0]]

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

Submission status: Time Limit Exceeded

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 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 = types.union({l[0]})

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

Submission status: Time Limit Exceeded

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