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: Runtime Error

Solution code:
import numpy as np
import heapq
import copy

N, K = list(map(int, input().split()))
t_list, d_list = [], []
for _ in range(N):
    t_, d_ = list(map(int, input().split()))
    t_list.append(t_)
    d_list.append(d_)
    
# Sort by delicoiousness.
c = zip(t_list, d_list)
c = sorted(c, key=lambda x: x[1])[-1::-1]
t_list, d_list = zip(*c)

first_sushi_set = []
nokori_set = []

iter_ = 0
for neta, delicious in zip(t_list, d_list):
    if iter_ < K:
        heapq.heappush(first_sushi_set, (delicious, neta))
    else:
        heapq.heappush(nokori_set, (-delicious, neta))
    iter_ += 1
    
first_neta_num = len(list(set(t_list[:K])))

def get_next_sushi_set(current_sushi, nokori_sushi):
    current_sushi_cp = copy.deepcopy(current_sushi)

    while(len(current_sushi) > 0):
        #種類数が減らないかつ一番美味しくない寿司を選ぶ
        remove_sushi = heapq.heappop(current_sushi)
        score, neta = remove_sushi
        if len(current_sushi) == 0:
            break
        if neta in np.array(current_sushi)[:, 1]:
            break
        else:
            pass
        
    current_sushi_cp.remove(remove_sushi)
    heapq.heappush(nokori_sushi, (-score, neta))
    nokori_sushi_cp = copy.deepcopy(nokori_sushi)

    while(len(nokori_sushi) > 0):
        #種類数が増えるかつ一番美味しい寿司を選ぶ
        add_sushi = heapq.heappop(nokori_sushi)
        negative_score, neta = add_sushi
        score = - negative_score
        if len(nokori_sushi) == 0:
            break
        if neta in np.array(current_sushi_cp)[:, 1]:
            pass
        else:
            break        
    heapq.heappush(current_sushi_cp, (score, neta))
    nokori_sushi_cp.remove(add_sushi)
            
    return current_sushi_cp, nokori_sushi_cp
        
current_sushi_set = first_sushi_set
nokori_sushi_set = nokori_set
score_list = []
for neta_num in range(first_neta_num, N+1, 1):
    score =  len(set(np.array(current_sushi_set)[:, 1]))**2 + sum(np.array(current_sushi_set)[:, 0])
    score_list.append(score)
    current_sushi_set, nokori_sushi_set = get_next_sushi_set(current_sushi_set, nokori_sushi_set)
print(max(score_list))

Submission status: Runtime Error

Solution code:
import numpy as np
import heapq
import copy

N, K = list(map(int, input().split()))
t_list, d_list = [], []
for _ in range(N):
    t_, d_ = list(map(int, input().split()))
    t_list.append(t_)
    d_list.append(d_)
    
# Sort by delicoiousness.
c = zip(t_list, d_list)
c = sorted(c, key=lambda x: x[1])[-1::-1]
t_list, d_list = zip(*c)

first_sushi_set = []
nokori_set = []

iter_ = 0
for neta, delicious in zip(t_list, d_list):
    if iter_ < K:
        heapq.heappush(first_sushi_set, (delicious, neta))
    else:
        heapq.heappush(nokori_set, (-delicious, neta))
    iter_ += 1
    
first_neta_num = len(list(set(t_list[:K])))

def get_next_sushi_set(current_sushi, nokori_sushi):
    current_sushi_cp = copy.deepcopy(current_sushi)

    while(len(current_sushi) > 0):
        #種類数が減らないかつ一番美味しくない寿司を選ぶ
        remove_sushi = heapq.heappop(current_sushi)
        score, neta = remove_sushi
        if len(current_sushi) == 0:
            break

        if neta in np.array(current_sushi)[:, 1]:
            break
        else:
            pass
        
    current_sushi_cp.remove(remove_sushi)
    heapq.heappush(nokori_sushi, (-score, neta))
    nokori_sushi_cp = copy.deepcopy(nokori_sushi)

    while(len(nokori_sushi) > 0):
        #種類数が増えるかつ一番美味しい寿司を選ぶ
        add_sushi = heapq.heappop(nokori_sushi)
        negative_score, neta = add_sushi
        score = - negative_score
        if len(nokori_sushi) == 0:
            break

        if neta in np.array(current_sushi_cp)[:, 1]:
            pass
        else:
            break        
    heapq.heappush(current_sushi_cp, (score, neta))
    nokori_sushi_cp.remove(add_sushi)
            
    return current_sushi_cp, nokori_sushi_cp
        
current_sushi_set = first_sushi_set
nokori_sushi_set = nokori_set
score_list = []
for neta_num in range(first_neta_num, N+1, 1):
    score =  len(set(np.array(current_sushi_set)[:, 1]))**2 + sum(np.array(current_sushi_set)[:, 0])
    score_list.append(score)
    current_sushi_set, nokori_sushi_set = get_next_sushi_set(current_sushi_set, nokori_sushi_set)
print(max(score_list))

Submission status: Runtime Error

Solution code:
from pulp import *

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

def get_duplicate_list(seq,index):
    seen = []
    return [x[index] for x in seq if not seen.append(x[index]) and seen.count(x[index]) == 2]

t=[0 for _ in range(N)]
d=[0 for _ in range(N)]
key=[i for i in range(N)]
sushi=[]
for i in range(N):
    t[i],d[i]=map(int,input().split())
    sushi.append([t[i],d[i]])

sorted_sushi=sorted(sushi,key= lambda x:x[1])
#print(sorted_sushi)
taberu=[]
for i in range(K):
    taberu.append(sorted_sushi.pop())

sorted_taberu = sorted ( taberu, key=lambda x: -x[ 1 ] )
maxscore=0
for i in range(N):
    # 重複要素をだす
    if len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu]) >maxscore:
        maxscore=len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu])
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sorted_taberu = sorted ( sorted_taberu, key=lambda x: -x[ 1 ],reverse=True )
    dp=get_duplicate_list(sorted_taberu,0)
    '''
    print('pop taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    if not dp:
        if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
            maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] )
        continue
    else:
        sb=False
        for i in range(len(sorted_taberu)):
            for s in dp:
                if sorted_taberu[i][0]==s:
                    sorted_sushi.append(sorted_taberu.pop(i))
                    sb=True
                    break
            if sb:
                break
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sb=False
    '''
    print('push taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    for i in range(len(sorted_sushi)):
        setted_sushi=set([x[0] for x in sorted_sushi])
        setted_taberu=set([x[0] for x in sorted_taberu])

        seki=setted_sushi & setted_taberu

        if not seki:
            sorted_taberu.append(sorted_sushi.pop())
            if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
                maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] )
                sb=True
                break

        for i in range(len(sorted_sushi)):
            for s in list(seki):
                if sorted_sushi[i][0]!=seki:
                    sb=True
                    sorted_taberu.append ( sorted_sushi.pop (i ) )
                    break
            if sb:
                break
        if sb:
            break




print (maxscore)
exit ( 0 )

Submission status: Runtime Error

Solution code:
from pulp import *

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

def get_duplicate_list(seq,index):
    seen = []
    return [x[index] for x in seq if not seen.append(x[index]) and seen.count(x[index]) == 2]

t=[0 for _ in range(N)]
d=[0 for _ in range(N)]
key=[i for i in range(N)]
sushi=[]
for i in range(N):
    t[i],d[i]=map(int,input().split())
    sushi.append([t[i],d[i]])

sorted_sushi=sorted(sushi,key= lambda x:x[1])
print(sorted_sushi)
taberu=[]
for i in range(K):
    taberu.append(sorted_sushi.pop())

sorted_taberu = sorted ( taberu, key=lambda x: -x[ 1 ] )
maxscore=0
for i in range(N):
    # 重複要素をだす
    if len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu]) >maxscore:
        maxscore=len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu])
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sorted_taberu = sorted ( sorted_taberu, key=lambda x: -x[ 1 ],reverse=True )
    dp=get_duplicate_list(sorted_taberu,0)
    '''
    print('pop taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    if not dp:
        if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
            maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] )
        continue
    else:
        sb=False
        for i in range(len(sorted_taberu)):
            for s in dp:
                if sorted_taberu[i][0]==s:
                    sorted_sushi.append(sorted_taberu.pop(i))
                    sb=True
                    break
            if sb:
                break
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sb=False
    '''
    print('push taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    for i in range(len(sorted_sushi)):
        setted_sushi=set([x[0] for x in sorted_sushi])
        setted_taberu=set([x[0] for x in sorted_taberu])

        seki=setted_sushi & setted_taberu

        if not seki:
            sorted_taberu.append(sorted_sushi.pop())
            if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
                maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] )
                sb=True
                break

        for i in range(len(sorted_sushi)):
            for s in list(seki):
                if sorted_sushi[i][0]!=seki:
                    sb=True
                    sorted_taberu.append ( sorted_sushi.pop (i ) )
                    break
            if sb:
                break
        if sb:
            break




print (maxscore)
exit ( 0 )

Submission status: Runtime Error

Solution code:
import numpy as np

n, k = map(int, input().split())
mx = {}
for _ in range(n):
    t, d  = map(int, input().split())
    if t not in mx:
        mx[t] = d
    elif mx[t] < d:
        mx[t] = d
    sushi.append((t,d))

# 美味しさでソート
taste_sorted_sushi = sorted(sushi, key=lambda x:x[1], reverse=True)

# 貪欲にK個選び分割する
selected_sushi = taste_sorted_sushi[:k]

# 貪欲なK個の種類の分布を調べる
# K個の中で二つ以上あるものとそうでないものに分ける
S = set()
q = []
for s in selected_sushi:
    if s[0] in S:
        q.append(s)
    else:
        S.add(s[0])

q = sorted(q, key=lambda x:x[1], reverse=True)
ans = sum(selected_sushi[:][1])+len(S)**2

S = set(selected_sushi[:][0])
for i in range(k,n):
    if taste_sorted_sushi[k] not in S:
        ans = max(ans, ans-q.pop()+taste_sorted_sushi[k]-len(S)**2 + (len(S)+1)**2
        S.add(taste_sorted_sushi[k][0])

print(ans)

Submission status: Runtime Error

Solution code:
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
import heapq
n,k=IL()
nl=dict()
for i in range(n+1):
    nl[i]=0
l=[]
ans=[]
f=[]
el=[]
heapq.heapify(el)
s=0
for i in range(n):
    t,d=IL()
    l.append([d,t])
l.sort(reverse=True)
for i in range(k):
    t=l.pop(0)
    el.heappush(t)
    nl[t[1]]+=1
    s+=t[0]
    if t[1] not in f:
        f.append(t[1])
u=len(f)
ans.append(s+u*u)
for i in range(u,k+1):
    if len(l)==0:
        break
    while True:
        t=el.heappop(0)
        if nl[t[1]]!=1:
            nl[t[1]]-=1
            s-=nl[t[0]]
            break
    while len(l)>0:
        t=l.pop(0)
        if nl[t[1]]==0:
            nl[t[1]]+=1
            s+=t[0]
            el.heappush(t)
            break
    ans.append(s+(u+i+1)**2)
    if len(l)==0:
        break
print(max(ans))

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = set()
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = sum(ans) + sum(aru) + s**2
print(ans, aru, s)

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans.append(nai[i][1])
        s += 1
        print(ans, aru, s)
    m = max(m, sum(ans) + sum(aru) + s**2)
print(m)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = []
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.append(anstmp[0])
    else:
        aru.append(anstmp[1])

#print(ans, spe, aru)
#print("nai", nai)
f.append(sum(ans) + sum(aru))
s = len(spe)
for i in range(len(nai)):
    if(nai[i][0] not in spe):
        aru.pop(-1)
        ans.append(nai[i][1])
        spe.append(nai[i][0])
        f.append(sum(ans) + sum(aru))
#print(f)
for i in range(len(f)):
    m = max(m, f[i] + (s + i)**2)
print(m)

Submission status: Runtime Error

Solution code:
import numpy as np
N, K = map(int,input().split())
td = np.empty((N,2),dtype=np.int64)
for i in range(N):
    td[i] = list(map(int,input().split()))

t = td[:,0]-1
d = td[:,1]
order1=np.argsort(d)[::-1]
mins = len(set(t[order1[:K]]))

best = -np.ones(N,dtype=np.int64)
for i in range(N):
    if best[t[i]] < 0 or d[i] > d[best[t[i]]]:
        best[t[i]]=i

bestmenu = np.array(list(set(best) - set([-1])),dtype=np.int64)
othermenu = np.array(list(set(range(N)) - set(bestmenu)),dtype=np.int64)

maxs = min(len(bestmenu),K)
ans = 0
bests=0
sorted_d_best=np.sort(d[bestmenu])
sorted_d_other=np.sort(d[othermenu])
bo
for s in range(mins,maxs+1):
    score = s*s + sorted_d_best[-s:].sum() + sorted_d_other[-(K-s):].sum()
    if score > ans:
        bests=s
        ans = score

print(ans)

Submission status: Runtime Error

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    if count < K:
        while 1:
            if count >= K :
                break
            if len(not_eat_list) <= 0:
                break
            i =not_eat_list.popleft()
            eat_list.appendleft(i)
            max_score += sushi[i][0]
            already_eat.add(sushi[i][1])
            count += 1



    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    if len(not_eat_list) > 0:
        for new_idx in not_eat_list:
            if len(eat_list) == 0 :
                break
            eated_idx = eat_list.pop()
            already_eat.remove(sushi[eated_idx][1])
            already_eat.add(sushi[new_idx][1])

            new_kind = len(already_eat)
            now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
            now_kind = new_kind
            max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]

td = sorted(td, key=lambda x: - x[1])

dict_td = set()
q = []
ans = 0
for i in range(k):
    if td[i][0] in dict_td:
        q.append(td[i])
    ans += td[i][1]
    dict_td.add(td[i][0])

ans += len(dict_td) ** 2

tmp_ans = ans
for i in range(k, n):
    if td[i][0] not in dict_td:
        minus = q.pop()
        tmp_ans += td[i][1] - minus[1]
        dict_len = len(dict_td)
        tmp_ans += (dict_len + 1) ** 2 - dict_len ** 2
        dict_td.add(td[i][0])
    ans = max(tmp_ans, ans)

print(ans)

Submission status: Runtime Error

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    if count < K:
        while 1:
            if count >= K :
                break
            if len(not_eat_list) <= 0:
                break
            i =not_eat_list.popleft()
            eat_list.appendleft(i)
            max_score += sushi[i][0]
            already_eat.add(sushi[i][1])
            count += 1



    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
from itertools import islice
import operator

N, K = map(int, input().split())
sushi = [tuple([i] + [int(x) for x in input().split()]) for i in range(N)]
sushi.sort(key=operator.itemgetter(2), reverse=True)

from collections import OrderedDict, deque, Counter
left = OrderedDict((x[0], x) for x in sushi)

base = 0
choosen = [None] * K
choosen_kind = set()
ci = 0
for x in sushi:
    i, t, _ = x
    if t not in choosen_kind:
        choosen[ci] = x
        ci += 1
        choosen_kind.add(t)
        del left[i]

choosen = choosen[:min(ci, K)]
base = sum(d for _,_,d in choosen)
kinds = len(choosen)

cands = islice(left.values(), K)
fixed = islice(cands, K - len(choosen))
base += sum(d for _,_,d in fixed)

score = base + kinds**2

for cand in cands:
    if cand[2] <= choosen[-1][2]:
        break

    base_n = base
    base_n -= choosen[-1][2]
    base_n += cand[2]
    kinds_n = kinds - 1
    
    score_n = base_n + kinds_n**2
    if score_n > score:
        choosen.pop()
        base, kinds, score = base_n, kinds_n, score_n
    else:
        break

print(score)

Submission status: Runtime Error

Solution code:
from operator import itemgetter

n, k = map(int, input().split())
td = sorted([list(map(int, input().split())) for i in range(n)], key=itemgetter(1), reverse=True)
best_sushi = []
best_syurui = set()
other = []
for i in range(n):
    if td[i][0] in best_syurui:
        other.append(td[i][1])
    else:
        best_sushi.append(td[i][1])
        best_syurui.add(td[i][0])
best_sushi_t = [0]
for i in range(len(best_sushi)):
    best_sushi_t.append(best_sushi[i] + best_sushi_t[i])
other_t = [0]
for i in range(len(other)):
    other_t.append(other[i] + other_t[i])
best = []
for i in range(1, k + 1):
    if k - i >= len(other_t):
        continue
    MP = best_sushi_t[i] + other_t[k - i]
    BP = len(best_sushi[:i])
    best.append(MP + BP ** 2)
print(max(best))

Submission status: Runtime Error

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = dequae([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    if count < K:
        while 1:
            if count >= K :
                break
            if len(not_eat_list) <= 0:
                break
            i =not_eat_list.popleft()
            eat_list.appendleft(i)
            max_score += sushi[i][0]
            already_eat.add(sushi[i][1])
        print max_score + len(already_eat) * len(already_eat)
        return


    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Runtime Error

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

for _ in range(N):
    t,d = map(int, input().split())
    T.append(t)
    D.append(d)

#全探索
#i層目、j選択中
def calc_max(i, j, L):
    if (N-1 <= i):
        res = 0
        if (j < K):
            res += D[i]
            L[T[i]-1] = 1
        l = sum(L)
        res += (l * l)
    else:
        if (j < K):
            L2 = L.copy()
            a = D[i]
            L2[T[i]-1] = 1
            a += calc_max(i+1, j+1, L2)
        else:
            a = 0
        if ((K - j) <= (N - i)):
          	b = calc_max(i+1, j, L)
        else:
            b = 0
        #print(i,j,a,b)
        res = max(a,b)
    return res

L = [0] * N
res = calc_max(0, 0, L)

print(res)

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter
from itertools import combinations

sys.setrecursionlimit(200000)


def input():
    return sys.stdin.readline()[:-1]


def ii(t: type = int):
    return t(input())


def il(t: type = int):
    return list(map(t, input().split()))


def imi(N: int, t: type = int):
    return [ii(t) for _ in range(N)]


def iml(N: int, t: type = int):
    return [il(t) for _ in range(N)]


def solve():
    TD = iml(N)
    N, K = il()
    ans = 0
    for td in combinations(TD, K):
        c = Counter([td[k][0] for k in range(K)])
        x = len(c)
        tmp = x ** 2 + sum([td[k][1] for k in range(K)])
        if ans < tmp:
            ans = tmp
    return ans


if __name__ == "__main__":
    print(solve())

Submission status: Runtime Error

Solution code:
import itertools
import operator

N, K = map(int, input().split())
sushi = [(i, *[int(x) for x in input().split()]) for i in range(N)]
sushi.sort(key=operator.itemgetter(2), reverse=True)

from collections import OrderedDict, deque, Counter
left = OrderedDict((x[0], x) for x in sushi)

base = 0
choosen = []
choosen_kind = set()
for x in sushi:
    i, t, _ = x
    if t not in choosen_kind:
        choosen.append(x)
        choosen_kind.add(t)
        del left[i]

choosen.sort(key=operator.itemgetter(2), reverse=True)
choosen = choosen[:K]
base = sum(d for _,_,d in choosen)
kinds = len(choosen)

tl = list(itertools.islice(left.values(), K))
fixed = tl[:K - len(choosen)]
cands = tl[K - len(choosen):]
base += sum(d for _,_,d in fixed)

score = base + kinds**2

if cands:
    comp = []
    while cands and cands[0][2] > choosen[-1][2]:
        base_n = base
        base_n -= choosen[-1][2]
        base_n += cands[0][2]
        kinds_n = kinds - 1
        
        score_n = base_n + kinds_n**2
        if score_n > score:
            choosen = choosen[:-1]
            fixed += cands[0]
            cands = cands[1:]
            base, kinds, score = base_n, kinds_n, score_n
        else:
            break
print(score)

Submission status: Runtime Error

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(eat_list)
    while 1:
        i =not_eat_list.popleft()
        if count >= K :
            break
        if len(not_eat_list) <= 0:
            break
        eat_list.appendleft(i)
        max_score += sushi[i][0]




    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Runtime Error

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)
#print(sushi)
num = 0
kind = set()
changable = []
for i in range(k):
	num += sushi[i][1]
	if sushi[i][0] in kind:
		changable.append(sushi[i][1])
	else:
		kind.add(sushi[i][0])

num += len(kind)**2
changable.reverse()
#print(changable)
ans = num

for i in range(k, n):
	if sushi[i][0] not in kind:
		num += sushi[i][1] + 2*len(kind)+1 - changable.pop()
		#print(num, kind)
		kind.add(sushi[i][0])
		ans = max(ans, num)

	if not changable:
		break

print(ans)

Submission status: Runtime Error

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)
#print(sushi)
num = 0
kind = set()
changable = []
for i in range(k):
	num += sushi[i][1]
	if sushi[i][0] in kind:
		changable.append(sushi[i][1])
	else:
		kind.add(sushi[i][0])

num += len(kind)**2
changable.reverse()
#print(changable)
ans = num

for i in range(k, n):
	if sushi[i][0] not in kind:
		num += sushi[i][1] + 2*len(kind)+1 - changable.pop()
		#print(num, kind)
		kind.add(sushi[i][0])
		ans = max(ans, num)

print(ans)