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:
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 = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    TD = np.array(TD)
    TD = TD[np.argsort(TD[:, 1])[::-1]]

    # print(TD)
    # 種類のカウントをもつ
    type_count = [0 for i in range(max(TD[:, 0]) + 1)]

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

    # print('type_count', type_count)
    base_point = np.sum(S[:, 1])
    type_bonus = len([i for i in type_count if i != 0])
    result = base_point + type_bonus**2
    # print('base_point, type_bonus, result', base_point, type_bonus**2, result)

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

    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) == 0:
            break
        if type_count[s[0]] >= 2:
            # print('s', s)
            type_count[s[0]] -= 1
            base_point -= s[1]
            #  S[len(S) - 1 - idx] = TD[remaining_items[0]]
            type_bonus += 1
            # print('base_point1', base_point)
            base_point += remaining_items[0]
            # print('base_point2', base_point)
            # print('type_bonus', type_bonus)
            del remaining_items[0]
            result = max(result, type_bonus**2 + base_point)

    print(result)
    
main()

Submission status: Time Limit Exceeded

Solution code:
N,K = list(map(int,input().split()))

TDlist = [list(map(int,input().split())) for x in range(N)]

TDlist.sort(key = lambda x: x[1],reverse=True)
kinds = []
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.append(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer
max_score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.append(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    max_score = max(score,answer)
    

answer = max(answer,max_score)

print(answer)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int, input().split())
netas_oishisas = []
for _ in range(n):
    netas_oishisas.append(tuple(map(int, input().split())))
netas_oishisas.sort(key = lambda no: no[1], reverse = True)
tabetaneta = []
second = []
greedyscore = 0
for i in range(k):
    neta, oishisa = netas_oishisas[i]
    if neta in tabetaneta:
        second.append(oishisa)
    else:
        tabetaneta.append(neta)
    greedyscore += oishisa
greedyscore += len(tabetaneta) ** 2
maxscore = greedyscore
currentscore = greedyscore
second.reverse()
secondpointer = 0
secondn = len(second)
for i in range(k, n):
    if secondpointer >= secondn:
        break
    neta, oishisa = netas_oishisas[i]
    if neta not in tabetaneta:
        currentscore = currentscore - second[secondpointer] + 2*len(tabetaneta) + 1 + oishisa
        maxscore = max(maxscore, currentscore)
        secondpointer += 1
        tabetaneta.append(neta)
print(maxscore)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int, input().split())
netas_oishisas = []
for _ in range(n):
    netas_oishisas.append(tuple(map(int, input().split())))
netas_oishisas.sort(key = lambda no: no[1], reverse = True)
tabetaneta = []
second = []
greedyscore = 0
for i in range(k):
    neta, oishisa = netas_oishisas[i]
    if neta in tabetaneta:
        second.append(oishisa)
    else:
        tabetaneta.append(neta)
    greedyscore += oishisa
greedyscore += len(tabetaneta) ** 2
maxscore = greedyscore
currentscore = greedyscore
for i in range(k, n):
    if second == []:
        break
    neta, oishisa = netas_oishisas[i]
    if neta not in tabetaneta:
        currentscore = currentscore - second[-1] + 2*len(tabetaneta) + 1 + oishisa
        maxscore = max(maxscore, currentscore)
        tabetaneta.append(neta)
print(maxscore)

Submission status: Time Limit Exceeded

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

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

TDlist.sort(key = lambda x: x[1],reverse=True)
kinds = []
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.append(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer
max_score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.append(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    max_score = max(score,answer)
    

answer = max(answer,max_score)

print(answer)

Submission status: Time Limit Exceeded

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

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

sorted_list = sorted(TDlist, key = lambda x: x[1], reverse=True)
kinds = []
seconds = []
answer = 0

for x in sorted_list[:K]:
    if not x[0] in kinds:
        kinds.append(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer
max_score = answer

for x in sorted_list[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.append(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    max_score = max(score,answer)
    

answer = max(answer,max_score)

print(answer)

Submission status: Time Limit Exceeded

Solution code:
def takeSecond(elem):
    return elem[1]

n,k = map(int,input().split())

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)

selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])
maxscore = score
front_pointer = k - 1
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for p in range(front_pointer,-1,-1):
        item = selected[p]
        flag = False
        for idx in range(p):
            if(selected[idx][0]==item[0]):
                score -= item[1]
                selected.remove(item)
                front_pointer = p - 1
                flag = True
                break
        if(flag):
            break

    #add the largest one that add kind
    for j in range(pointer,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            pointer = j + 1            
            score += 2*startKindNum + 1 + sushi[j][1]
            if(score>maxscore):
                maxscore = score
            startKindNum += 1
            break
print(maxscore)

Submission status: Time Limit Exceeded

Solution code:
def takeSecond(elem):
    return elem[1]

n,k = map(int,input().split())

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)

selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])

scorelist = []
scorelist.append(score)
front_pointer = k - 1
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for p in range(front_pointer,-1,-1):
        item = selected[p]
        flag = False
        for idx in range(p):
            if(selected[idx][0]==item[0]):
                score -= item[1]
                selected.remove(item)
                front_pointer = p - 1
                flag = True
                break
        if(flag):
            break

    #add the largest one that add kind
    for j in range(pointer,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            pointer = j + 1            
            score += 2*startKindNum + 1 + sushi[j][1]
            startKindNum += 1
            scorelist.append(score)
            break
print(max(scorelist))

Submission status: Time Limit Exceeded

Solution code:
def takeSecond(elem):
    return elem[1]

n,k = map(int,input().split())

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)

selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])

scorelist = []
scorelist.append(score)
front_pointer = k - 1
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for p in range(front_pointer,-1,-1):
        item = selected[p]
        if([elem for elem in selected if(elem[0]==item[0] and elem is not item)]):
            score -= item[1]
            selected.remove(item)
            front_pointer = p - 1
            break

    #add the largest one that add kind
    for j in range(pointer,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            pointer = j + 1            
            score += 2*startKindNum + 1 + sushi[j][1]
            startKindNum += 1
            scorelist.append(score)
            break
print(max(scorelist))

Submission status: Time Limit Exceeded

Solution code:
def takeSecond(elem):
    return elem[1]

n,k = map(int,input().split())

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)

selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])

scorelist = []
scorelist.append(score)
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for item in selected[::-1]:
        if([elem for elem in selected if(elem[0]==item[0] and elem is not item)]):
            score -= item[1]
            selected.remove(item)
            break

    #add the largest one that add kind
    for j in range(k,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            
            score += 2*startKindNum + 1 + sushi[j][1]
            startKindNum += 1
            scorelist.append(score)
            break
print(max(scorelist))

Submission status: Time Limit Exceeded

Solution code:
def takeSecond(elem):
    return elem[1]

n,k = map(int,input().split())

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)

selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])

scorelist = []
scorelist.append(score)
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for item in selected[::-1]:
        if([elem for elem in selected if(elem[0]==item[0] and elem is not item)]):
            selected.remove(item)
            break

    #add the largest one that add kind
    for j in range(k,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            kindNum = len(kindset)
            score = kindNum*kindNum + sum([item[1] for item in selected])
            scorelist.append(score)
            break
print(max(scorelist))

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi.sort(key=lambda x: x[1], reverse=True)

maximums = []
remains = []
kinds = set()
for t, d in sushi[:K]:
    if t in kinds:
        remains.append(d)
    else:
        maximums.append(d)
    kinds.add(t)
score = sum(maximums) + sum(remains) + len(kinds)**2

for t, d in sushi[K:]:
    if not remains:
        break
    if t in kinds:
        continue
    kinds.add(t)
    remains.pop()
    maximums.append(d)
    test = sum(maximums) + sum(remains) + len(kinds)**2
    score = max(score, test)

print(score)

Submission status: Time Limit Exceeded

Solution code:
from heapq import heappop, heappush, heapify

N, K = map(int, input().split())
sushi = []
for _ in range(N):
    t, d = map(int, input().split())
    sushi.append((-d, t))
heapify(sushi)

maxs = []
rems = []
kinds = set()
for _ in range(K):
    d, t = heappop(sushi)
    if t in kinds:
        heappush(rems, -d)
    else:
        maxs.append(-d)
        kinds.add(t)
ans = sum(maxs) + sum(rems) + len(kinds)**2

while True:
    if not sushi or not rems:
        break
    d, t = heappop(sushi)
    if t in kinds:
        continue
    else:
        heappop(rems)
        maxs.append(-d)
        kinds.add(t)
        ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from heapq import heappop, heappush

N, K = map(int, input().split())
sushi = []
for _ in range(N):
    t, d = map(int, input().split())
    heappush(sushi, (-d, t))

maxs = []
rems = []
kinds = set()
for _ in range(K):
    d, t = heappop(sushi)
    if t in kinds:
        heappush(rems, -d)
    else:
        maxs.append(-d)
        kinds.add(t)
ans = sum(maxs) + sum(rems) + len(kinds)**2

while True:
    if not sushi or not rems:
        break
    d, t = heappop(sushi)
    if t in kinds:
        continue
    else:
        heappop(rems)
        maxs.append(-d)
        kinds.add(t)
        ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from itertools import groupby
from heapq import heapify, heappop

N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]

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

maxs = []
rems = []
kinds = set()
for k, g in groupby(sorted(sushi[:K], key=lambda x: x[0]), key=lambda x: x[0]):
    d = list(list(zip(*g))[1])
    maxs.append(d.pop(0))
    rems.extend(d)
    kinds.add(k)
heapify(rems)
ans = sum(maxs) + sum(rems) + len(kinds)**2

for k, g in groupby(sushi[K:], key=lambda x: x[0]):
    if not rems:
        break
    if k in kinds:
        continue
    d = list(list(zip(*g))[1])
    maxs.append(d.pop(0))
    heappop(rems)
    kinds.add(k)
    ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from itertools import groupby
from heapq import heapify, heappop

N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]

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

maxs = []
rems = []
kinds = set()
for k, g in groupby(sorted(sushi[:K], key=lambda x: x[0]), key=lambda x: x[0]):
    d = list(list(zip(*g))[1])
    maxs.append(d.pop(0))
    rems.extend(d)
    kinds.add(k)
heapify(rems)
ans = sum(maxs) + sum(rems) + len(kinds)**2

for t, d in sushi[K:]:
    if not rems:
        break
    if t in kinds:
        continue
    maxs.append(d)
    heappop(rems)
    kinds.add(t)
    ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from  operator import itemgetter

# 満足ポイントを求める関数
def f(a, b):
    return len(a) ** 2 + sum(a) + sum(b)

n, k = map(int, input().split())

sushi = []
for i in range(n):
    sushi.append(list(map(int, input().split())))

sushi.sort(key=itemgetter(1), reverse=True)

s = []
s_change = []
types = []
# 集合Sの生成
for i in range(k):
    if sushi[i][0] not in types:
        s.append(sushi[i])
        types.append(sushi[i][0])
    else:
        s_change.append(sushi[i])

s_change.sort(key=itemgetter(1))
#print(s)
#print(s_change)

m = []
# 交換用の集合の生成
for i in range(k, n):
    if sushi[i][0] not in types:
        m.append(sushi[i])
        types.append(sushi[i][0])

#print(m)
#print([l[1] for l in s])
g = []
g.append(f([l[1] for l in s], [l[1] for l in s_change]))

# 集合gを求める
for i in range(min(len(s_change), len(m))):
    s.append(m[i])
    s_change.pop(0)
    g.append(f([l[1] for l in s], [l[1] for l in s_change]))

print(max(g))

Submission status: Time Limit Exceeded

Solution code:
from  operator import itemgetter
import numpy as np

n, k = map(int, input().split())

sushi = []
for i in range(n):
    sushi.append(list(map(int, input().split())))

#print(sushi)

sushi.sort(key=itemgetter(1), reverse=True)

#print(sushi)

sushi = np.array(sushi)

#print(sushi)

total = len(set(sushi[:,0]))

#print(total)

initial = len(set(sushi[:k, 0]))

#print(initial)

g = []

#print(initial, total)
for i in range(initial, min(total, k)+1):
    check = [0 for i in range(n)]
    value = i * i
    types = []
    s = 0
    for j in range(i):
        while True:
            if sushi[s, 0] not in types:
                types.append(sushi[s, 0])
                value += sushi[s, 1]
                check[s] = 1
                s += 1
                break
            s += 1

    s = 0

    for j in range(k-i):
        while True:
            if check[s] == 0:
                value += sushi[s, 1]
                s += 1
                break
            s += 1
    
    g.append(value)

#print(g)
print(max(g))

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,raw_input().split())

data = []
for i in range(N):
    t,d = map(int,raw_input().split())
    data.append([d,t])

data.sort(reverse = True)

sub = []
leave = []
syurui = []
count = 0
point = [0,0]
for i in range(N):
    if data[i][1] not in syurui:
        syurui.append(data[i][1])
        data[i] = [data[i][0],1]
        if count >= K:
            sub.append(data[i])
    else:
        data[i] = [data[i][0],0]
        if count < K:
            leave.append(data[i])
    if count < K:
        if data[i][1] == 1:
            point[1] += 1
        point[0] += data[i][0]
    count += 1


ans = point[0]+point[1]*point[1]

while len(sub) != 0 and len(leave) != 0:
    point[0] += sub[0][0]-leave[-1][0]
    point[1] += 1
    del leave[-1],sub[0]
    ans = max(ans,point[0]+point[1]*point[1])

print ans

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,raw_input().split())

data = []
for i in range(N):
    t,d = map(int,raw_input().split())
    data.append([d,t])

data.sort(reverse = True)

sub = []
eat = []
leave = []
syurui = []
count = 0
for i in range(N):
    if data[i][1] not in syurui:
        syurui.append(data[i][1])
        data[i] = [data[i][0],1]
        if count >= K:
            sub.append(data[i])
    else:
        data[i] = [data[i][0],0]
        if count < K:
            leave.append(data[i])
    if count < K:
        eat.append(data[i])
    count += 1

leave.reverse()


point = [0,0]
count = 0
for i in eat:
    if i[1] == 1:
        point[1] += 1
    point[0] += i[0]
ans = point[0]+point[1]*point[1]

while len(sub) != 0 and len(leave) != 0:
    eat.remove(leave[0])
    eat.append(sub[0])
    point[0] += sub[0][0]-leave[0][0]
    point[1] += 1
    del leave[0],sub[0]
    ans = max(ans,point[0]+point[1]*point[1])

print ans

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,raw_input().split())


def point(eat):
    point = 0
    count = 0
    for i in eat:
        if i[2] == 1:
            syurui.append(i[2])
            count += 1
        point += i[0]
    point += count*count
    return point


data = []
for i in range(N):
    t,d = map(int,raw_input().split())
    data.append([d,t])

data.sort(reverse = True)

syurui = []
sub = []
leave = []
count = 0
for i in range(N):
    if data[i][1] not in syurui:
        syurui.append(data[i][1])
        data[i] = [data[i][0],data[i][1],1]
        if count >= K:
            sub.append(data[i])
    else:
        data[i] = [data[i][0],data[i][1],0]
        if count < K:
            leave.append(data[i])
    count += 1

eat = []
for i in range(K):
    eat.append(data[i])
eat.reverse()
leave.reverse()

ans = point(eat)
for i in range(len(eat)):
    if len(sub) != 0 and len(leave) != 0:
        eat.remove(leave[0])
        eat.append(sub[0])
        del leave[0]
        del sub[0]
        ans = max(ans,point(eat))

print ans

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict
import heapq


def get_first_try(K, ts, ds):
    eats = []
    non_eats = []
    for i in range(len(ts)):
        t = ts[i]
        d = ds[i]
        heapq.heappush(eats, (d, t))
        if len(eats) > K:
            non_eats.append(heapq.heappop(eats))
    return [(t, d) for d,t in eats], [(t, d) for d,t in non_eats]

def count_kinds(eats):
    kinds = defaultdict(int)
    for e in eats:
        kinds[e[0]] += 1
    return kinds

def calc_satisfaction(eats, eat_kinds):
    s = 0
    for e in eats:
        s += e[1]
    return s + calc_bonus(eat_kinds)


def calc_bonus(kinds):
    x = 0
    for k in kinds:
        if kinds[k] > 0:
            x += 1
    return x**2


def solve(N, K, ts, ds):
    eats, non_eats = get_first_try(K, ts, ds)
    eat_kinds = count_kinds(eats)
    satisfaction = calc_satisfaction(eats, eat_kinds)
    sorted_eats = sorted(eats, key=lambda x: x[1])
    sorted_non_eats = sorted(non_eats, key=lambda x: x[1], reverse=True)

    max_satisfaction = satisfaction

    i = 0
    j = 0
    while i < len(sorted_eats) and j < len(sorted_non_eats):
        t, d = sorted_eats[i]
        if eat_kinds[t] == 1:
            i += 1
            continue
        t2, d2 = sorted_non_eats[j]
        if eat_kinds[t2] > 0:
            j += 1
            continue
        prev_bonus = calc_bonus(eat_kinds)
        eat_kinds[t] -= 1
        eat_kinds[t2] += 1
        current_bonus = calc_bonus(eat_kinds)
        satisfaction = satisfaction - d + d2 - prev_bonus + current_bonus
        if max_satisfaction < satisfaction:
            max_satisfaction = satisfaction
        i += 1
        j += 1

    return max_satisfaction


N, K = list(map(int, input().split(' ')))

ts = []
ds = []
for i in range(N):
    t, d = list(map(int, input().split(' ')))
    ts.append(t)
    ds.append(d)

print(solve(N, K, ts, ds))

Submission status: Time Limit Exceeded

Solution code:
import copy
from collections import namedtuple

Sushi = namedtuple('Sushi', ('t', 'd'))

n, k = map(int, input().split(' '))
all_sushi = []
for i in range(n):
    t, d = map(int, input().split(' '))
    all_sushi.append(Sushi(t=t, d=d))

all_sushi.sort(key=lambda e: e.d, reverse=True)

all_t = set(map(lambda e:e.t, all_sushi))
all_t_len = len(all_t)

#
# Ignore types and select more delicious order.
#
selected_sushi = all_sushi[:k]
addable_sushi = all_sushi[k:]
cur_d = sum(map(lambda e:e.d, selected_sushi))
all_t = set(map(lambda e:e.t, all_sushi))
cur_t = set(map(lambda e:e.t, selected_sushi))
cur_t_len = len(cur_t)
cur_sc = cur_d + cur_t_len * cur_t_len
best_sc = cur_sc
#print(f'cur_sc={cur_sc}')

#
# Prepare to add sushi.
#
add_t = all_t - cur_t
add_sushi = []
for t in add_t:
    s = next(filter(lambda e:e.t == t, addable_sushi), None)
    if s is not None:
        add_sushi.append(s)
add_sushi = sorted(add_sushi, key=lambda e:e.d, reverse=True)
#print(f'add_sushi={add_sushi}')

#
# Try to decrement number of types.
#
removable_sushi = copy.copy(selected_sushi)
for t in cur_t:
    r = next(filter(lambda e:e.t == t, selected_sushi), None)
    removable_sushi.remove(r)

#print(f'removable_sushi={removable_sushi}')
try:
    while True:
        r = removable_sushi.pop()
        a = add_sushi.pop(0)
        if (r.d - a.d >= all_t_len * all_t_len - cur_t_len * cur_t_len):
            break
        cur_sc += a.d - r.d
        cur_sc -= cur_t_len * cur_t_len
        cur_t_len += 1
        cur_sc += cur_t_len * cur_t_len
        #print(f'cur_sc={cur_sc}')
        if best_sc < cur_sc:
            best_sc = cur_sc
except IndexError:
    pass

print(best_sc)

Submission status: Time Limit Exceeded

Solution code:
import copy
from collections import namedtuple

def first_true(iterable, default=None, pred=None):
    """Returns the first true value in the iterable.

    If no true value is found, returns *default*

    If *pred* is not None, returns the first item
    for which pred(item) is true.

    """
    # first_true([a,b,c], x) --> a or b or c or x
    # first_true([a,b], x, f) --> a if f(a) else b if f(b) else x
    return next(filter(pred, iterable), default)

Sushi = namedtuple('Sushi', ('t', 'd'))

n, k = map(int, input().split(' '))
all_sushi = []
for i in range(n):
    t, d = map(int, input().split(' '))
    all_sushi.append(Sushi(t=t, d=d))

all_sushi.sort(key=lambda e: e.d, reverse=True)

all_t = set(map(lambda e:e.t, all_sushi))
all_t_len = len(all_t)

#
# Ignore types and select more delicious order.
#
selected_sushi = all_sushi[:k]
addable_sushi = all_sushi[k:]
cur_d = sum(map(lambda e:e.d, selected_sushi))
all_t = set(map(lambda e:e.t, all_sushi))
cur_t = set(map(lambda e:e.t, selected_sushi))
cur_t_len = len(cur_t)
cur_sc = cur_d + cur_t_len * cur_t_len
best_sc = cur_sc
#print(f'cur_sc={cur_sc}')

#
# Prepare to add sushi.
#
add_t = all_t - cur_t
add_sushi = []
for t in add_t:
    s = next(filter(lambda e:e.t == t, addable_sushi), None)
    if s is not None:
        add_sushi.append(s)
add_sushi = sorted(add_sushi, key=lambda e:e.d, reverse=True)
#print(f'add_sushi={add_sushi}')

#
# Try to decrement number of types.
#
removable_sushi = copy.copy(selected_sushi)
for t in cur_t:
    r = next(filter(lambda e:e.t == t, selected_sushi), None)
    removable_sushi.remove(r)

#print(f'removable_sushi={removable_sushi}')
while (len(removable_sushi) > 0) and (len(add_sushi) > 0):
    r = removable_sushi[-1]
    removable_sushi = removable_sushi[:-1]
    a = add_sushi[0]
    add_sushi = add_sushi[1:]
    if (r.d - a.d >= all_t_len * all_t_len - cur_t_len * cur_t_len):
        break
    cur_sc -= r.d
    cur_sc += a.d
    cur_sc -= cur_t_len * cur_t_len
    cur_t_len += 1
    cur_sc += cur_t_len * cur_t_len
    #print(f'cur_sc={cur_sc}')
    if best_sc < cur_sc:
        best_sc = cur_sc

print(best_sc)

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
typ = []
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.append(item[0])
        item[0] = 1
        lstA.append(item[1])
ans = 0
neta = 0
for i in range(K):
    ans += sushi[i][1]
    if sushi[i][0] == 1:
        neta += 1
ans += neta ** 2
C = sum(lstA[:neta])
D = sum(lstB[:K - neta])
for i in range(neta + 1, min(K, len(lstA)) + 1):
    C += lstA[i - 1]
    D -= lstB[K - i]
    tmp = C + D + i ** 2
    ans = max(ans, tmp)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
from pprint import pprint
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
typ = []
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.append(item[0])
        item[0] = 1
        lstA.append(item[1])
ans = 0
neta = 0
for i in range(K):
    ans += sushi[i][1]
    if sushi[i][0] == 1:
        neta += 1
ans += neta ** 2
C = sum(lstA[:neta])
D = sum(lstB[:K - neta])
for i in range(neta + 1, min(K, len(lstA)) + 1):
    C += lstA[i - 1]
    D -= lstB[K - i]
    tmp = C + D + i ** 2
    ans = max(ans, tmp)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import itertools

n,k = map(int,input().split())

td = []
for i in range(n):
    td_tmp = list(map(int,input().split()))
    td.append(td_tmp)

# 貪欲で

td_comb = list(itertools.combinations(td, k))

ans = 0

for i in td_comb:
    tmp_ans = 0
    cnt = 0
    t_lst = []
    for j in i:
        tmp_ans += j[1]
        if j[0] not in t_lst:
            cnt += 1
            t_lst.append(j[0])
    tmp_ans += cnt**2
    ans = tmp_ans if tmp_ans>ans else ans

print(ans)

Submission status: Time Limit Exceeded

Solution code:
#ABC116D-VariousSushi.py

from sys import stdin
from numpy import sum

input = stdin.readline

N, K = [int(n) for n in input().split()]
INPUT = [[None for i in range(2)] for j in range(N)]
List1, List2 = [], []
kind1 = []

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

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

for i in range(N):
    if INPUT[i][0] not in kind1:
            List1.append([INPUT[i][1], 1])
            kind1.append(INPUT[i][0])
    else: List1.append([INPUT[i][1], 0])

List2 = List1[K:]
if List2 != []:
    List2.sort(key =lambda x : (x[1],x[0]))
    del List1[K:]
    lenList2 = len(List2)
else: lenList2 = 0

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

sum = sum(List1, axis = 0)

Score = sum[0] + sum[1]**2
Test = Score
kindnumber = sum[1]

for i in range(lenList2):
    if List2[-1][1] == 0 or List1[-1][1] == 1:
        break
    Test = Test - List1.pop()[0] + List2.pop()[0] + 2 * kindnumber + 1
    kindnumber = kindnumber + 1
    if Test >= Score:
        Score = Test

print(int(Score))