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 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)
    for i in not_eat_list:
        if count >= K :
            break
        if len(not_eat_list) <= 0:
            break
        eat_list.appendleft(not_eat_list.popleft())




    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:
import sys

sys.setrecursionlimit(100000)

N, K = map(long, raw_input().split())
S = [map(long, raw_input().split()) for _ in range(N)]
memo = {set([]):{0:0}}
def manzoku(i, k, kinds, basis):
  if memo[kinds][basis]:
    return memo[kinds][basis]
  if k == K:
    return basis+len(kinds)**2
  elif i >= N:
    return 0
  else:
    t = kinds.copy()
    t.add(S[i][0])
    return max(manzoku(i+1, k+1, t,basis+S[i][1]),
      manzoku(i+1, k,kinds,basis))
print manzoku(0, 0, set([]), 0)

Submission status: Runtime Error

Solution code:
from operator import itemgetter

def calc(sushis):
	global k
	manzoku_sushi = sum(list(map(lambda sushi: sushi[1], sushis[:k])))
	syurui = len(set(list(map(lambda sushi: sushi[0], sushis[:k]))))

	# print(sushis)
	# print(set(sushis))
	# print(list(map(lambda sushi: sushi[1], sushis)))
	# print(set(list(map(lambda sushi: sushi[1], sushis))))
	# print(manzoku_sushi, syurui)
	return manzoku_sushi + syurui*syurui

def calcb(sushis):
	global k
	manzoku_sushi, syurui = 0 , 0
	for index in range(k):
		if index + 1 <= len(sushis) and sushis[index][1] == sushis[index+1][1]:
			continue
		else:
			manzoku_sushi += sushis[index][1]
			syurui += 1

	return manzoku_sushi + syurui*syurui



n, k = map(int, input().split())
sushis = [tuple(map(int, input().split())) for _ in range(n)]

total_1 = calc( sorted(sushis, key=itemgetter(1), reverse=True) )
total_2 = calc( sorted(sushis, key=itemgetter(0), reverse=False) )
total_3 = calcb(sorted(sushis, key=itemgetter(1), reverse=True) )
total_4 = calcb( sorted(sushis, key=itemgetter(0), reverse=False) )


# print(total_1, total_2, total_3, total_4))

print(max(total_1, total_2, total_3, total_4))

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)
    for i in not_eat_list:
        if count >= K :
            break
        eat_list.appendleft(not_eat_list.popleft())




    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(long, raw_input().split())
S = [map(long, raw_input().split()) for _ in range(N)]
def manzoku(i, k, kinds, basis):
  if k == K:
    return basis+len(kinds)**2
  elif i >= N:
    return 0
  else:
    t = kinds.copy()
    t.add(S[i][0])
    return max(manzoku(i+1, k+1, t,basis+S[i][1]),
      manzoku(i+1, k,kinds,basis))
print manzoku(0, 0, set([]), 0)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10001)

def solve(N, K, tds):
    memos = {}
    tds.sort()

    def dfs(n, k, ee, neta_count):
        if (n, k, ee, neta_count) in memos:
            return memos[(n, k, ee, neta_count)]

        if n == N or k == 0:
            memos[(n, k, ee, neta_count)] = neta_count * neta_count
            return memos[(n, k, ee, neta_count)]

        if ee == tds[n][0]:
            eaten = dfs(n + 1, k - 1, tds[n][0], neta_count) + tds[n][1]
        else:
            eaten = dfs(n + 1, k - 1, tds[n][0], neta_count + 1) + tds[n][1]

        uneaten = dfs(n + 1, k, ee, neta_count)

        memos[(n, k, ee, neta_count)] = max(eaten, uneaten)
        return memos[(n, k, ee, neta_count)]

    return dfs(0, K, -1, 0)

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    tds = [tuple(map(int, input().split(" "))) for _ in range(N)]
    print(solve(N, K, tds))

Submission status: Runtime Error

Solution code:
#!usr/bin/env python3
def LI(): return list(map(int, input().split()))
def II(): return int(input())
def LS(): return input().split()
def S(): return input()
def LIR(n): return [LI() for i in range(n)]
mod = 1000000007

#A
"""
a,b,c = LI()
s = a+b+c
s /= 2
ans = (s*(s-a)*(s-b)*(s-c))**(1/2)
print(int(ans))
"""

#B
"""
s = int(input())
m = 1
lis = [True for i in range(1000001)]

while lis[s]:
    lis[s] = False
    if s % 2 == 1:
        s = 3*s+1
    else:
        s //= 2
    m += 1
print(m)
"""

#C
"""
n = II()
h = LI()
ans = 0
while sum(h) != 0:
    l = 0
    while h[l] == 0: l += 1
    r = l
    while r < n:
        if h[r] == 0: break
        h[r] -= 1
        r += 1
    ans += 1
print(ans)
"""

#D
n,k = LI()
lis = LIR(n)
lis.sort(key = lambda x:x[1])
lis = lis[::-1]
num = [False for i in range(n)]
count = [0 for i in range(n)]
for i in range(n):
    lis[i][0] -= 1
ans = 0
t = 0
l = []
for i in range(k):
    ans += lis[i][1]
    count[lis[i][0]] += 1
    if not num[lis[i][0]]:
        t += 1
    num[lis[i][0]] = True
for i in range(k):
    if count[i] != 1:
        l.append(lis[i])
l = l[::-1]

a = [0 for i in range(len(l))]
for i in range(len(l)):
    a[i] = l[i][1]
for i in range(len(l)-1):
    a[i+1] += a[i]
li = []

for i in range(n):
    if not num[lis[i][0]]:
        li.append(lis[i][1])
        num[lis[i][0]] = True
b = [0 for i in range(len(li))]
for i in range(len(li)):
    b[i] = li[i]
for i in range(len(li)-1):
    b[i+1] += b[i]
d = [[ans,t] for i in range(min(len(l),len(li))+1)]

for i in range(len(d)-1):
    d[i+1][0] += b[i] - a[i]
    d[i+1][1] = min(k,t+i+1)
ans = 0
for i in range(len(d)+1):
    ans = max(ans, d[i][0]+d[i][1]**2)
print(ans)
#E

#F

#G

#H

#I

#J

#K

#L

#M

#N

#O

#P

#Q

#R

#S

#T

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 21:13:02 2019

@author: Yamazaki Kenichi
"""

N,K = map(int,input().split())
A =[list(map(int,input().split())) for i in range(N)]

A.sort(key=lambda x: x[1])
kinds = 0
t = 0
for i in range(N):
    if t != A[i][0]:
        kinds += 1
        t = A[i][0]
c = K - kinds if K - kinds >=0 else 0
d = K-c
temp_t = pow(K-c,2)
i = N-1
counta =0
temp_d=0
a = [0 for i in range(N+1)]
b,b2 =[],[]
while counta < K:
    if c-a[A[i][0]] >= 0:
        temp_d += A[i][1]
        c = c-1 if a[A[i][0]] >= 1 else c
        a[A[i][0]] += 1
        counta+=1
        b2.append(A[i][1])
    else:
        b.append(A[i][1])
    i -=1
ans = temp_t+temp_d
#print(ans)
for i in range(len(b)):
    if b[i]-b2[-1-i]>pow(d,2)-pow(d-i-1,2):
        ans = ans +b[i]-b2[-1-i]-(pow(d,2)-pow(d-i-1,2))
    else:
        break
print(ans)

Submission status: Runtime Error

Solution code:
def solve(N, K, tds):
    memos = {}
    tds.sort()

    def dfs(n, k, ee, neta_count):
        if (n, k, ee, neta_count) in memos:
            return memos[(n, k, ee, neta_count)]

        if n == N or k == 0:
            memos[(n, k, ee, neta_count)] = neta_count ** 2
            return memos[(n, k, ee, neta_count)]

        if ee == tds[n][0]:
            eaten = dfs(n + 1, k - 1, tds[n][0], neta_count) + tds[n][1]
        else:
            eaten = dfs(n + 1, k - 1, tds[n][0], neta_count + 1) + tds[n][1]

        uneaten = dfs(n + 1, k, ee, neta_count)

        memos[(n, k, ee, neta_count)] = max(eaten, uneaten)
        return memos[(n, k, ee, neta_count)]

    return dfs(0, K, -1, 0)

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    tds = [tuple(map(int, input().split(" "))) for _ in range(N)]
    print(solve(N, K, tds))

Submission status: Runtime Error

Solution code:
#!usr/bin/env python3
def LI(): return list(map(int, input().split()))
def II(): return int(input())
def LS(): return input().split()
def S(): return input()
def LIR(n): return [LI() for i in range(n)]
mod = 1000000007

#A
"""
a,b,c = LI()
s = a+b+c
s /= 2
ans = (s*(s-a)*(s-b)*(s-c))**(1/2)
print(int(ans))
"""

#B
"""
s = int(input())
m = 1
lis = [True for i in range(1000001)]

while lis[s]:
    lis[s] = False
    if s % 2 == 1:
        s = 3*s+1
    else:
        s //= 2
    m += 1
print(m)
"""

#C
"""
n = II()
h = LI()
ans = 0
while sum(h) != 0:
    l = 0
    while h[l] == 0: l += 1
    r = l
    while r < n:
        if h[r] == 0: break
        h[r] -= 1
        r += 1
    ans += 1
print(ans)
"""

#D
n,k = LI()
lis = LIR(n)
lis.sort(key = lambda x:x[1])
lis = lis[::-1]
num = [False for i in range(n)]
count = [0 for i in range(n)]
for i in range(n):
    lis[i][0] -= 1
ans = 0
t = 0
l = []
for i in range(k):
    ans += lis[i][1]
    count[lis[i][0]] += 1
    if not num[lis[i][0]]:
        t += 1
    num[lis[i][0]] = True
for i in range(k):
    if count[i] != 1:
        l.append(lis[i])
l = l[::-1]

a = [0 for i in range(len(l))]
for i in range(len(l)):
    a[i] = l[i][1]
for i in range(len(l)-1):
    a[i+1] += a[i]
li = []

for i in range(n):
    if not num[lis[i][0]]:
        li.append(lis[i][1])
        num[lis[i][0]] = True
b = [0 for i in range(len(li))]
for i in range(len(li)):
    b[i] = li[i]
for i in range(len(li)-1):
    b[i+1] += b[i]
d = [[ans,t] for i in range(min(len(l),len(li))+1)]

for i in range(min(len(li),len(l))):
    d[i+1][0] += b[i] - a[i]
    d[i+1][1] = min(k,t+i+1)
ans = 0
for i in range(len(li)+1):
    ans = max(ans, d[i][0]+d[i][1]**2)
print(ans)
#E

#F

#G

#H

#I

#J

#K

#L

#M

#N

#O

#P

#Q

#R

#S

#T

Submission status: Runtime Error

Solution code:
import math
import itertools

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

dt = []
neta = []
bass_score = 0
netamax = []
for i in range(N):
    if i == 0:
        tmp1,tmp2 = map(int, input().split())
        netamax.append(tmp2)
        dt.append([tmp2,tmp1])
    else:
        tmp1,tmp2 = map(int, input().split())
        if dt[i-1] != tmp1:
            netamax.append(tmp2)
        dt.append([tmp2,tmp1])


DT = dt.copy()
DT.sort(reverse=True)

for i in range(K):
    bass_score += DT[i][0]
    neta.append(DT[i][1])

bonus_score = len(set(neta))**2
test1 = bass_score + bonus_score
#----------------------------
netamax.sort(reversr=True)
bass_score = sum(netamax[:K])

test2 = bass_score + K*K

print(max(test1,test2))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
S = [list(map(int, input().split())) for n in range(N)]
T = {0 : [0]}
maxT = {}
for s in S:
    t = T.setdefault(s[0], [])
    t.append(s[1])

for t, l in T.items():
    l.sort(reverse=True)
    maxT[t] = l[0]


def max_s():
    L = [max(l) for l in T.values()]
    L.sort(reverse=True)
    if len(L) >= K:
        return sum(L[:K]) + K * K
    else:
        res = sum(L) + len(L) * len(L)
        R = K - len(L)
        L2 = [l[1:] for t, l in T.items() if t != 0]
        L2.sort(reverse=True)
        res += sum(L2[:R])
        return res

def get_max_t(T):
    res_t = -1
    res = 0
    for t, v in maxT.items():
        if res < v:
            res = v
            res_t = t
    return res_t, res


max_k = max_s()
res = 0
s = 0
for k in range(K):
    t, d = get_max_t(maxT)
    d_0 = T[0][0]

    v = d - s ** 2 + (s + 1) ** 2
    if v < d_0:
        res += d_0
        T[0].pop(0)
    else:
        res += v
        s += 1
        T[t].pop(0)
        T[0].extend(T[t])
        T[0].sort(reverse=True)
        maxT.pop(t)


print(max(res, max_k))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
items = []
neta = {}
for _ in range(N):
    t, d = tuple(map(int, input().split()))
    items.append((t, d))

    if t not in neta:
        neta[t] = d
    else:
        tmp = neta[t]
        if d > tmp:
            neta[t] = d

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

ans1 = 0
select1 = set()
for i in range(K):
    t, d = items[i]
    select1.add(t)
    ans1 += d
ans1 += len(select1) ** 2

ans2 = 0
i = 1
select2 = set()
cnt = 0
for k, v in neta:
    cnt += 1
    if v + cnt ** 2 - (cnt - 1) ** 2 < items[i][0]:
        ans2 += items[i][0]
        cnt += 1
    else:
        select2.add(k)
        ans2 += v
    if cnt == K:
        ans2 += cnt ** 2
        print(max(ans1, ans2))
        exit()
ans2 += cnt ** 2

remain = K - cnt
sel = set()
i += 1
while remain != 0:
    t, d = items[i]
    if t not in sel:
        i += 1
        sel.add(t)

    else:
        ans2 += d
        sel.add(t)
        remain -= 1
        i += 1
print(max(ans1, ans2))

Submission status: Runtime Error

Solution code:
#!/Users/takahiro.kida/git/atcorder/venv/bin/python

import numpy
import scipy

from collections import defaultdict
from pprint import pprint

N, K = [int(x) for x in input().split(" ")]
sushis = [] # type, delicious, eat?
for n in range(N):
    row = [int(x) for x in input().split(" ")] + [False]
    sushis.append(row)



type_of_max = {}  # type: [index of sushis, deli]
for i in range(N):
    sushi = sushis[i]
    if not sushi[0] in type_of_max:
        type_of_max[sushi[0]] = [i, sushi[1]]
        continue
    if type_of_max[sushi[0]][1] < sushi[1]:
        type_of_max[sushi[0]] = [i, sushi[1]]
        continue


vali = 0
eat = []
for type in type_of_max:
    index = type_of_max[type][0]
    eat.append(sushis[index])
    sushis[index][2] = True
    vali += 1

eat = sorted(eat, key=lambda sushi: sushi[1], reverse=True)
dish = []
for sushi in sushis:
    if not sushi[2]:
        dish.append(sushi)
dish = sorted(dish, key=lambda sushi: sushi[1], reverse=True)

if len(eat) > K:
    for i in range(K, len(eat) - 1):
        sushi = eat.pop()
        dish.append(sushi)
        vali += -1

if len(eat) < K:
    for i in range(0, K - len(eat) -1):
        sushi = dish.pop()
        eat.append(sushi)

dish = sorted(dish, key=lambda sushi: sushi[1], reverse=True)
eat = sorted(eat, key=lambda sushi: sushi[1], reverse=True)

while True:
    d_diff = dish[0][1] - eat[-1][1]
    t_diff = vali * vali - (vali - 1) * (vali - 1)
    if (d_diff > t_diff):
        eat.pop()
        eat.insert(0, dish.pop(0))
        vali += -1
    else:
        break

deli = [x[1] for x in eat]
print(vali*vali + sum(deli))

Submission status: Runtime Error

Solution code:
import itertools
import cython
n,k = map(int,input().split())
l = []
for i in range(n):
    l.append(list(map(int,input().split())))
comb = list(itertools.combinations(l,k))
def calc_score(comb):
    t = set()
    cdef int d = 0
    for i in comb:
        t.add(i[0])
        d += i[1]
    return len(t) ** 2 + d
print(max([calc_score(i) for i in comb]))

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)
    for i in not_eat_list:
        if count >= K:
            break
        eat_list.appendleft(not_eat_list.popleft())




    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        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())

S=[[]for i in range(10**5+1)]
for i in range(N):
    a,b=map(int,input().split())
    S[a].append(b)
S.sort(reverse=True)
#print(S[:10])

k=0
o=0
E=[]
for i in range(len(S)):
    if S[i]==[]:
        break
    else:
        S[i].sort(reverse=True)
        k+=1
        o+=S[i][0]
        E.append(S[i][0])
        del S[i][0]
E.sort()
#print(S)
S.sort(reverse=True)
A=[]
if k>K:
    for i in range(k-K):
        A.append(E[i])
        o-=E[i]
    for i in range(k-K):
        del E[i]
        k-=1
#print(A)
for i in range(len(S)):
    if S[i]==[]:
        break
    else:
        for j in range(len(S[i])):
            A.append(S[i][j])
A.sort(reverse=True)
#print(A)
if K-k>0:
    for i in range(K-k):
        o+=A[i]
        del A[i]
        #print(i)
    
    #print(A)
    #print(E)
for i in range(min(K,len(A))):
    if k*k+o<=(k-1)**2+o-E[i]+A[i]:
        k=(k-1)
        o=o-E[i]+A[i]
    else:
        break
print(k*k+o)

Submission status: Runtime Error

Solution code:
N,K=list(map(int,input().split()))
import numpy as np
td=np.array([list(map(int,input().split())) for i in range(N)])
T=td[:,0]
D=td[:,1]
dp=np.zeros(N-K)
sumk=np.sum(D[:K])+np.unique(T[:K]).shape[0]**2
dp[0]=sumk
for i in range(1,N-K):
    for j in range(K):
        max=0
        tmp=np.sum(np.append(D[:j],D[j+1]))+D[K+i]+np.unique(np.append(np.append(D[:j],D[j+1:K])),D[k+i]).shape[0]**2
        if max<tmp:
            max=tmp
    dp[i]=max
print(int(dp[N-K-1]))

Submission status: Runtime Error

Solution code:
N, K = [int(i) for i in input().split()]
td = [tuple(int(i) for i in input().split()) for i in range(N)]

def f(i, eat_t, base_vp, td_):
    if not td_:
        return base_vp + len(eat_t) ** 2
    elif i >= K:
        return f(i+1, eat_t, base_vp, td_[1:])
    else:
        t, d = td_[0]
        return max(
            f(i+1, eat_t, base_vp, td_[1:]),
            f(i+1, eat_t | set([t]), base_vp + d, td_[1:])
        )
print(f(0, set(), 0, td))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
neta = [[] for i in range(N+1)]
for i in range(N):
    t, d = map(int, input().split())
    neta[t].append(d)
for i in range(N):
    neta[i].sort(reverse = True)
neta = [i for i in neta if i != []]
neta.sort(reverse = True)
kinds = len(neta)
maxneta = 0
for i in range(len(neta)):
    maxneta = max(maxneta, len(neta[i]))
anslist = []
kinds_max = min(kinds, K)
x = K
for i in range(len(neta)):
    x -= len(neta[i])
    if x <= 0:
        kinds_min = i+1
        break
firstnetasum = [0]
for i in range(kinds):
    firstnetasum.append(firstnetasum[i] + neta[i][0])
for i in range(kinds_min, kinds_max + 1):
    ans = firstnetasum[i]
    remain = []
    remaincount = K-i
    for j in range(i):
        try:
            remain += neta[j][1:]
            remaincount -= len(a)
            if remaincount < 0:
                break
        except IndexError:
            pass
    remain.sort(reverse = True)
    ans += sum(remain[:K-i])
    ans += i**2
    anslist.append(ans)

print(max(anslist))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
S = [list(map(int, input().split())) for n in range(N)]
T = {0 : [0]}
maxT = {}
for s in S:
    t = T.setdefault(s[0], [])
    t.append(s[1])

for t, l in T.items():
    l.sort(reverse=True)
    maxT[t] = l[0]


def max_s():
    L = [max(l) for l in T.values()]
    L.sort(reverse=True)
    return sum(L[:K]) + K * K


def get_max_t(T):
    res_t = -1
    res = 0
    for t, v in T.items():
        if res < v:
            res = v
            res_t = t
    return res_t, res


max_k = max_s()
res = 0
s = 0
for k in range(K):
    t, d = get_max_t(maxT)
    d_0 = T[0][0]

    v = d - s ** 2 + (s + 1) ** 2
    if v < d_0:
        res += d_0
        T[0].pop(0)
    else:
        res += v
        s += 1
        T[t].pop(0)
        T[0].extend(T[t])
        T[0].sort(reverse=True)
        maxT.pop(t)


print(max(res, max_k))

Submission status: Runtime Error

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

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

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

ans = 0
used = [a[0][0]]

ans += a[0][1]
a.pop(0)

for i in range(K - 1):
    sp = len(set(used))
    for _ in a:
        if _[0] not in used:
            cand = _

    naive = (ans + a[0][1])
    abnormal = (ans + (sp + 1) * (sp + 1)) + _[1]
    if naive > abnormal:
        ans += a[0][1]
        used.append(a[0][0])
        a.pop(0)
    else:
        ans += cand[1]
        a.remove(cand)
        used.append(cand[0])

sp = len(set(used))
ans += sp * sp
print(ans)

Submission status: Runtime Error

Solution code:
n,k=map(int,input().split())
s=[]
d=[]
p=[]
t=[]
for i in range(n):
    a,b=map(int,input().split())
    p.append(b)
    if a not in s:
        s.append(a)
        d.append([b])
    else:
        d[s.index(a)].append(b)    
for i in range(len(s)):
    d[i].sort()
    t.append(d[i][-1])
t.sort()
p.sort()
r=[]
q=0
g=0
for i in range(k):
    if i<=len(t)+1:
        r.append(t[-1])
        q+=t[-1]
        g+=1
        p.pop(p.index(t[-1]))
        t.pop(-1)
    else:
        q+=sum(p[n-k+len(t):])
        del p[n-k+len(t):]
r.sort()
while r[0]+2*g-1<p[-1]:
    q+=p[-1]-r[0]
    g-=1
    p.pop(-1)
    r.pop(0)
print(g**2+q)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
n,k = map(int, input().split())
o = {}

for _ in range(n):
    t,d = map(int, input().split())
    if t not in o:
        o[t] = []
    o[t].append(d)

tl = list(o.keys())
tl.sort(reverse=True,key=lambda x:o[x])
# print(tl)
for t in tl:
    o[t].sort(reverse=True)
# print(o)
res = -1
sel = {}
for t in tl:
    if len(o[t])<=k:
        sel[t] = len(o[t])
        k -= sel[t]
    else:
        sel[t] = k
        k = 0
    if k==0:
        break

# print(sel)

def calc(sel):
    res = 0
    cnt = 0
    for t in sel:
        if sel[t]>0:
            cnt += 1
        for i in range(sel[t]):
            res += o[t][i]
    return res+cnt*cnt

idx = None
for i in range(len(tl)-1):
    if tl[i+1] not in sel:
        idx = i
        break

if idx is None:
    print(calc(sel))
    exit()

# print(idx)

res = calc(sel)
for i in range(idx+1,len(tl)):
    while sel[tl[idx]]==1:
        idx -= 1

    sel[tl[idx]] -= 1
    sel[i] = 1
    # res2 = res + o[i][0] - o[tl[idx]][-1]
    # res = max(res, res2)
    res = max(res, calc(sel))

print(res)

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 = []
    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

    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        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:
import itertools
n,k = map(int,input().split())
su = [list(map(int,input().split())) for i in range(n)]
sushi = list(itertools.combinations(su,k))
p = []
t = []
d = []
for i in range(n):
  for j in range(k):
    t.append(sushi[i-1][j-1][0])
    d.append(sushi[i-1][j-1][1])
  T = len(list(set(t)))
  ans = sum(d) + T * T
  p.append(ans)
  t.clear()
  d.clear()
print(max(p))

Submission status: Runtime Error

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

S=[[]for i in range(10**5+1)]
for i in range(N):
    a,b=map(int,input().split())
    S[a].append(b)
S.sort(reverse=True)
#print(S[:10])

k=0
o=0
E=[]
for i in range(len(S)):
    if S[i]==[]:
        break
    else:
        S[i].sort(reverse=True)
        k+=1
        o+=S[i][0]
        E.append(S[i][0])
        del S[i][0]
E.sort()
#print(S)
S.sort(reverse=True)
A=[]
if k>K:
    for i in range(k-K):
        A.append(E[i])
        o-=E[i]
    for i in range(k-K):
        del E[i]
        k-=1
#print(A)
for i in range(len(S)):
    if S[i]==[]:
        break
    else:
        for j in range(len(S[i])):
            A.append(S[i][j])
A.sort(reverse=True)
#print(A)
if K-k>0:
    for i in range(K-k):
        o+=A[i]
        #print(i)
    del A[:i+1]
    #print(A)
    #print(E)
for i in range(min(K,len(A))):
    if k*k+o<=(k-1)**2+o-E[i]+A[i]:
        k=(k-1)
        o=o-E[i]+A[i]
    else:
        break
print(k*k+o)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
n,k = map(int, input().split())
o = {}

for _ in range(n):
    t,d = map(int, input().split())
    if t not in o:
        o[t] = []
    o[t].append(d)

tl = list(o.keys())
tl.sort(reverse=True,key=lambda x:o[x])
# print(tl)
for t in tl:
    o[t].sort(reverse=True)
# print(o)
res = -1
sel = {}
for t in tl:
    if len(o[t])<=k:
        sel[t] = len(o[t])
        k -= sel[t]
    else:
        sel[t] = k
        k = 0
    if k==0:
        break

# print(sel)

def calc(sel):
    res = 0
    cnt = 0
    for t in sel:
        if sel[t]>0:
            cnt += 1
        for i in range(sel[t]):
            res += o[t][i]
    return res+cnt*cnt

idx = None
for i in range(len(tl)-1):
    if tl[i+1] not in sel:
        idx = i
        break

if idx is None:
    print(calc(sel))
    exit()

# print(idx)

res = calc(sel)
for i in range(idx+1,len(tl)):
    while sel[tl[idx]]==1:
        idx -= 1

    sel[tl[idx]] -= 1
    sel[i] = 1
    res2 = res + o[i][0] - o[tl[idx]][-1]
    res = max(res, res2)
    # res = max(res, calc(sel))

print(res)