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:
# -*- coding: utf-8 -*-
"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d

"""
import sys

from collections import Counter

def solve(N, K, sushi):
    sushi.sort(key=lambda x: x[1], reverse=True)
    pick = sushi[:K]
    rem = sushi[K:][::-1]
    all_type = set([t for t, d in sushi])
    pick_type = set([t for t, d in pick])

    cnt = Counter()
    for t, d in pick:
        cnt[t] += 1

    manzoku = sum([d for t, d in pick])
    ans = [manzoku + len(pick_type)**2]
    while pick_type != all_type:
        t, d = rem.pop()
        if t not in pick:
            pick_type.add(t)
            tt, dd = pick.pop()
            if cnt[tt] > 1:
                cnt[tt] -= 1
                manzoku = manzoku - dd + d
                ans.append(manzoku+ len(pick_type)**2)
    return max(ans)



def main(args):
    N, K = map(int, input().split())
    sushi = [[int(i) for i in input().split()] for _ in range(N)]
    ans = solve(N, K, sushi)
    print(ans)


if __name__ == '__main__':
    main(sys.argv[1:])

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dt = sorted(zip(d, t), reverse=True)

    q0 = []

    acc = 0
    ts = 0
    p = [-1] * (N+1)
    for dd, tt in dt[:K]:
        acc += dd
        v = min(p[tt], dd)
        p[tt] = max(p[tt], dd)
        if 0 < v:
            heappush(q0, (v, tt))
        else:
            ts += 1
    
    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

    mmm = acc + ts * ts
    while 0 < len(q1):
        d1, t1 = heappop(q1)
        if 0 < p[t1]:
            continue
        acc -= d1
        p[t1] = -d1
        ts += 1

        d0, _ = heappop(q0)
        acc -= d0
        mmm = max(mmm, acc + ts * ts)

    return mmm


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    result = solve(N, K, t, d)
    print(result)

if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10**7)
INF = 10 ** 18
MOD = 10 ** 9 + 7
from functools import partial, reduce
from operator import mul
prod = partial(reduce, mul)
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 II(): return int(sys.stdin.readline())
def SI(): return input()
from collections import Counter

fibo = [1, 1]
for i in range(2, 9):
    fibo.append(fibo[i - 1] + fibo[i - 2])

def main():
    N, K = LI()
    TD = Counter()
    for _ in range(N):
        t, d = LI()
        TD[t] = max(TD[t], d)
    ds = list(TD.values())
    ds.sort(reverse=True)
    tmp = ds[0] * K
    ans = tmp + 1
    for i in range(1, K):
        tmp = tmp + ds[i] - ds[0]
        ans = max(ans, tmp + ((i+1)**2))
    return ans

print(main())

Submission status: Runtime Error

Solution code:
from collections import Counter
import sys
INF = 10 ** 18
def LI(): return [int(x) for x in sys.stdin.readline().split()]
def MI(): return map(int, sys.stdin.readline().split())


N, K = MI()
# 美味しさ降順寿司
sushi = sorted([LI() for _ in range(N)], key=lambda x:x[1], reverse=True)
# 美味しい寿司を貪欲に取る
delicious_sushi = sushi[:K]
# 余った寿司
soudemonai_sushi = sushi[K:]

syurui_counter = Counter()
deliciousness = 0

f = [-INF for _ in range(K+1)]

for x in delicious_sushi:
    syurui_counter[x[0]] += 1
    deliciousness += x[1]

f[len(syurui_counter)] = deliciousness

# print(delicious_sushi)
# print(soudemonai_sushi)
# print(syurui_counter)
# print("==========")


delicious_sushi = delicious_sushi[::-1]

i = 0
j = 0
while len(syurui_counter) < K and 0 < len(soudemonai_sushi):
    # print("len(syurui_counter) = " + str(len(syurui_counter)))
    while i < K:
        # print("i = " + str(i))
        if syurui_counter[delicious_sushi[i][0]] >= 2:
            # print(delicious_sushi[i])
            break
        i += 1
    while j < N - K:
        # print("j = " + str(j))
        if syurui_counter[soudemonai_sushi[j][0]] == 0:
            # print(soudemonai_sushi[j])
            break
        j += 1
    deliciousness -= delicious_sushi[i][1]
    deliciousness += soudemonai_sushi[j][1]
    syurui_counter[delicious_sushi[i][0]] -= 1
    syurui_counter[soudemonai_sushi[j][0]] += 1
    # print(syurui_counter)
    f[len(syurui_counter)] = deliciousness
    i += 1
    j += 1
    # print("==========")

# print(f)
for i in range(len(f)):
    f[i] += (i * i)
# print(f)
print(max(f))

Submission status: Runtime Error

Solution code:
from collections import Counter
import sys
INF = 10 ** 18
def LI(): return [int(x) for x in sys.stdin.readline().split()]
def MI(): return map(int, sys.stdin.readline().split())


N, K = MI()
# 美味しさ降順寿司
sushi = sorted([LI() for _ in range(N)], key=lambda x:x[1], reverse=True)
# 美味しい寿司を貪欲に取る
delicious_sushi = sushi[:K]
# 余った寿司
soudemonai_sushi = sushi[K:]

syurui_counter = Counter()
deliciousness = 0

f = [-INF for _ in range(K+1)]

for x in delicious_sushi:
    syurui_counter[x[0]] += 1
    deliciousness += x[1]

f[len(syurui_counter)] = deliciousness

# print(delicious_sushi)
# print(soudemonai_sushi)
# print(syurui_counter)
# print("==========")


delicious_sushi = delicious_sushi[::-1]

i = 0
j = 0
while len(syurui_counter) < K:
    # print("len(syurui_counter) = " + str(len(syurui_counter)))
    while i < K:
        # print("i = " + str(i))
        if syurui_counter[delicious_sushi[i][0]] >= 2:
            # print(delicious_sushi[i])
            break
        i += 1
    while j < N - K:
        # print("j = " + str(j))
        if syurui_counter[soudemonai_sushi[j][0]] == 0:
            # print(soudemonai_sushi[j])
            break
        j += 1
    deliciousness -= delicious_sushi[i][1]
    deliciousness += soudemonai_sushi[j][1]
    syurui_counter[delicious_sushi[i][0]] -= 1
    syurui_counter[soudemonai_sushi[j][0]] += 1
    # print(syurui_counter)
    f[len(syurui_counter)] = deliciousness
    i += 1
    j += 1
    # print("==========")

# print(f)
for i in range(len(f)):
    f[i] += (i * i)
# print(f)
print(max(f))

Submission status: Runtime Error

Solution code:
# 漸化式的に解く
# 一般性を失わない範囲をちゃんと考える
# 2種類以上あるものは、同じ種類のものとして差し支えない

from collections import defaultdict
N, K = map(int,input().split())
dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]

# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = [0] * (N+1)

cnt = 0
oisii = []
for t,d in pairs[:K]: #K個上から順番に取ってみる
    kinds[t] += 1
    cnt += d
    # print(t,d)
    oisii.append([t,d])

point = sum([d for t,d in oisii]) + sum([1 for kind in kinds if kind!=0])**2
ans = point
for t,d in pairs[K:]:
    """
    置き換えようとしていたけど
    その必要はなく、
    ほしいのはmax値だけなので、配列そのものを更新し続ける必要はない
    """
    # if kinds[t] == 1:
    #     continue
    # else:
    #     index = K
    #     while True:
    #         t_min, d_min = oisii[index-1]
    #         # 2個以上ある種類のなかで最小値を求める
    #         if kinds[t_min] >= 2:
    #             break
    #         index -= 1
    #         if index == 0:
    #             break
    # if index == 0:
    #     break
    if kinds[t] >= 1:
        continue
    # とりあえず最小だったoisiiを置き換えていく
    t_min, d_min = oisii.pop()
    n = sum([1 for kind in kinds if kind!=0])
    # とりあえず種類を追加
    kinds[t] += 1
    kinds[t_min] -= 1
    point = point - (d_min + n**2)  + (d + (n+1)**2)
    # print(point)
    ans = max(point,ans)
# 集計
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
td = sorted([list(map(int, input().split())) for i in range(n)], reverse=True, key=lambda x: x[1])

type = set()
L = []
Sum = 0
for x in td[:k]:
    Sum += x[1]
    if x[0] not in type:
        type.add(x[0])
    else:
        L.append(x[1])

L = L[::-1]
type_cnt = len(type)
ans = Sum + type_cnt ** 2
for x in td[k:]:
    if x[0] not in type:
        type.add(x[0])
        type_cnt += 1
        Sum = Sum - L.pop(0) + x[1]
        ans = max(ans, Sum + type_cnt ** 2)
print(ans)

Submission status: Runtime Error

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

# 美味しさ順
base_var = set([_td[0] for _td in td[:K]])
score_d = sum([_td[1] for _td in td[:K]]) + len(base_var)**2

score_list = [score_d]

candidate_list = []
for _td in td[K:]:
    if _td[0] not in base_var:
        candidate_list.append(_td)
        base_var.add(_td[0])


td_dict = {}
for _td in td[:K]:
    if _td[0] not in td_dict:
        td_dict[_td[0]] = [_td[1]]
    else:
        td_dict[_td[0]].append(_td[1])
pop_list = []

for k, v in td_dict.items():
    if len(v) == 1:
        continue
    v.sort(reverse=True)
    for _v in v[1:]:
        pop_list.append([k, _v])
pop_list.sort(key=lambda x: x[1])

current = td[:K]

for i in range(len(candidate_list)):
    idx = current.index(pop_list[i])
    current.pop(idx)
    current.insert(0, candidate_list[i])
    var = set([_td[0] for _td in current])
    score = sum([_td[1] for _td in current]) + len(var) ** 2
    score_list.append(score)

print(max(score_list))

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys
from collections import defaultdict
from itertools import accumulate
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    sushi = defaultdict([])
    for i in range(N):
        sushi[t[i]].append(d[i])

    A = [] # 各ネタでいちばんおいしい寿司のおいしさ
    B = [] # 各ネタで2番目以降においしい寿司のおいしさリスト
    
    for v in sushi.values():
        v.sort(reverse=True)
        A.append(v[0])
        B.extend(v[1:])
    A.sort(reverse=True)
    B.sort(reverse=True)

    X = [0] + [x for x in accumulate(A)] # X[i]: Aからうまい順にi個選んだ時のおいしさ合計
    Y = [0] + [y for y in accumulate(B)] # Y[j]: Bからうまい順にj個選んだ時のおいしさ合計

    res = 0
    for i in range(1, min(len(X), K + 1)):
        if K - i < len(Y):
            res = max(res, X[i] + Y[K - i] + i ** 2) 
            # B[:i]に含まれるネタが、A[i+1:]に含まれていた場合、実際のネタ数はiより多いが、
            # その場合はiを多くとったほうが、より美味しさが大きくなる。
            # そのためres[i]が答えになりえないため、無視して良い。
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    X = [(d[i], t[i]) for i in range(N)]
    X.sort(key=lambda x: x[1], reverse=True)

    Q = [] # うまい順にK個
    M = dict() # ネタごとに一番うまいやつ
    for i in range(N):
        if d[i] > M.get(t[i], 0):
            M[t[i]] = d[i]
        if len(Q) < K:
            heapq.heappush(Q, (d[i], t[i]))
        else:
            dq, tq = heapq.heappop(Q)
            if d[i] > dq:
                dq, tq = d[i], t[i]
            heapq.heappush(Q, (dq, tq))
    t_all = M.keys() # ネタ総数
    
    M = [(dm, tm) for tm, dm in M.items()]
    Q.sort(reverse=True)
    M.sort(reverse=True)

    T = [[0, set()]] # T[i=0~K]: うまい方からi個選び、のこりをMから選んだ時のおいしさと種類
    for i in range(K):
        t = T[-1][0] + Q[i][0]
        s = T[-1][1] | set([Q[i][1]])
        T.append([t, s])

    for i in range(K-1, -1, -1):
        if len(M) == 0:
            break
        while True:
            dm, tm = M.pop(0)
            if tm in T[i][1]:
                continue
            break
        else:
            break
        T[i][1] = T[i+1][1] | set([tm])
        T[i][0] = T[i+1][0] + dm - Q[i][0]

    res = 0
    for i in range(K+1):
        total = T[i][0] + len(T[i][1]) ** 2
        if total > res:
            res = total
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
from collections import defaultdict
from itertools import accumulate

input = sys.stdin.readline

N, K = map(int, input().split())
sushi = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    sushi[t].append(d)

X, Y = [], []
for v in sushi.values():
    v = sorted(v, reverse=True)
    X.append(v[0])
    if len(v) > 1:
        Y += v[1:]
X = sorted(X, reverse=True)

X = [0] + [x for x in accumulate(X)]
Y = [0] + [x for x in accumulate(Y)]

ans = 0
for i in range(1, K + 1):
    j = K - i
    if j < len(Y):
        ans = max(ans, X[i] + Y[j] + i * i)
print(ans)

Submission status: Runtime Error

Solution code:
n,K=map(int,input().split())
s_list=[]
for i in range(n):
    s_list.append([])
for i in range(n):
    t,d=map(int,input().split())
    s_list[t-1].append(d)
s_list=[i for i in s_list if not i==[]]
n_list=[]
for i in s_list:
    i.sort(reverse=True)
    n_list.append(len(i))
dp={}
def f(i,v,k):
    if (i,v,k) in dp:
        return dp[(i,v,k)]
    else:
        if i==len(n_list):
            if k==0:
                return 0
            else:
                return -10**9
        elif k==0:
            return 0
        else:
            Max=f(i+1,v,k)
            for j in range(0,n_list[i]):
                if k-j-1>=0:
                    m=f(i+1,v+1,k-j-1)+2*v+1+sum(s_list[i][0:j+1])
                    if m>Max:
                        Max=m
            dp[(i,v,k)]=Max
            return(Max)
print(f(0,0,K))

Submission status: Runtime Error

Solution code:
from operator import itemgetter

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

neta_list = [[] for i in range(N)]
max_dict = {}
for i in range(N):
  t,d = map(int, input().split())
  neta_list[i] = [d,t]
  if t not in max_dict.keys():
    max_dict[t] = d
  else:
    max_dict[t] = max(d,max_dict[t])
neta_list.sort(reverse=True)
kind_max = len(max_dict)
max_index = {}
for kind in max_dict.keys():
  max_index[kind] = neta_list.index([max_dict[kind], kind])
max_list = list(max_dict.items())
max_list = sorted(max_list, key=itemgetter(1), reverse=True)

ans = 0
used_flag = [0]*(N+1)
kind = 0
for k in range(K):
  if used_flag[neta_list[0][1]] == 0:
    ans += neta_list[0][0]
    kind += 1
    used_flag[neta_list[0][1]] = 1
    neta_list.pop(0)
    max_list.pop(0)
  else:
    used_max = neta_list[0][0]
    unused_max = max_list[0][1]
    if unused_max + (kind_max*kind_max - kind*kind) >= used_max:
      ans += unused_max
      kind += 1
      used_flag[max_list[0][0]] = 1
      neta_list.pop(max_index[max_list[0][0]]-k)
      max_list.pop(0)
    else:
      ans += used_max
      neta_list.pop(0)

print(ans + kind*kind)

Submission status: Runtime Error

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

ty = set()
im = []
mu = []
an = []

for t,d in sushi:
	if t in ty:
		if K > 0:
			mu.append(d)
			K -= 1
	else:
		ty.add(t)
		im.append(d)
		if K > 0:
			K -= 1
		else:
			mu.pop()

	an.append(len(ty)**2 + sum(im) + sum(mu))

print(max(an))

Submission status: Runtime Error

Solution code:
N,K=map(int, input().split())
td=[list(map(int,input().split())) for i in range(N)]
best=[0]*(N+1)
other=[]
from heapq import heappop,heappush
td.sort(key=lambda x: x[1],reverse=True)
shurui=0
for i in range(K):
    if best[td[i][0]]==0:
        best[td[i][0]]=td[i][1]
        shurui+=1
    elif best[td[i][0]]<td[i][1]:
        heappush(other,best[td[i][0]])
        best[td[i][0]]=td[i][1]
    else:
        heappush(other,td[i][1])

ans=sum(best)+sum(other)+shurui**2

for i in range(K,N):
    if best[td[i][0]]==0:
        best[td[i][0]]=td[i][1]
        shurui+=1
        heappop(other)
        ans=max(ans,sum(best)+sum(other)+shurui**2)

print(ans)

Submission status: Runtime Error

Solution code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r = []
    variety = set([])
    ans = 0
    for _c in choice:
        ans += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    ans += len(variety)**2
    i = k
    while len(variety) < min(max_variety, k):
        while td[i][0] in variety:
            i += 1
        variety.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Runtime Error

Solution code:
N,K=map(int,input().split())
td=[]
for i in range(N):
    t,d=map(int,input().split())
    td.append([t,d])
td.sort(key=lambda x:x[1],reverse=True)
A=td[:K]
B=[]
s2=set()
for i,j in td[K:]:
    if not(i in s2):
        s2.add(i)
        B.append([i,j])
s1=set()
n=0
C=[]
sum=0
for i,j in A:
    if not(i in s1):
        s1.add(i)
        n+=1
        sum+=j
    else:
        C.append([i,j])
        sum+=j
sum+=n**2
B.reverse()
D=[]
for i,j in B:
    if not(i in s1):
        D.append([i,j])
list=[sum]
while(D):
    i,j=D.pop()
    list.append(list[-1]+2*n+1+(j-A[-1][1]))
    n+=1
    A.pop()
    if n==K:
        break
print(max(list))

Submission status: Runtime Error

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

###dが大きい順にソート
td.sort(key = lambda x: x[1],reverse=True)

###各tでdが最大のものを格納
fix_eated_d=[]
###各tでdが最大でないものを格納。tの種類を増やす際の削除候補
change_eated_d=[]

eated_t = set()
d_point = 0

###まずdを最大化
for t,d in td[:k]:
  ###すでにtが選ばれているなら最大のdではないのでchangeへ
  if t in eated_t:
    change_eated_d.append(d)
    
  ###まだtが選ばれていないなら最大のdだからfixへ
  else:
    fix_eated_d.append(d)
    eated_t.add(t)
  d_point += d
point = len(eated_t)**2 + d_point

###tの種類を増やす
for t,d in td[k:]:
  ###すでにtが選ばれているならskip
  if t in eated_t:
    continue
    
  ###まだtが選ばれていないなら追加、change中dが最小のものを削除
  eated_t.add(t)
  d_point += d - change_eated_d.pop()
  point = max(point, len(eated_t)**2 + d_point)
  
  ###削除候補がなくなれば終了
  if not change_eated_d:
    break
  
print(point)

Submission status: Runtime Error

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

###dが大きい順にソート
td.sort(key = lambda x: x[0],reverse=True)
td.sort(key = lambda x: x[1],reverse=True)

###なるべくtの種類を多いなかで、dを最大化する
###まず各tでdが最大のものを1つずつ選ぶ
eated_kind = []
eated_kind_num = 0
eated = []
for t,d in td:
  if t not in eated_kind:
    eated_kind.append(t)
    eated_kind_num += 1
    eated.append([t,d])
    
    ###tがk種類存在したら、残りのtを扱うことはない
    ###またdも最大化されているため、これ以降のtdを使うことはない
    if eated_kind_num == k:
      td = td[:i+1]
      break

not_eated = [i for i in td if i not in eated]

###tがk種類存在しなかったら、k個になるまでtに関係なくdを最大化する
if eated_kind_num < k:
  eated += not_eated[:k-eated_kind_num]
  not_eated = not_eated[k-eated_kind_num:]
  #print(eated)
  eated.sort(key = lambda x: x[1],reverse=True)
  
t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point
###種類の寿司でdの最大化をする
for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  eated = eated[:-1]
  point = max(point,point+t_point_sub+d_point_sub)
  eated_kind_num -= 1
  if eated_kind_num == 1:
    break
#print(eated)

print(point)

Submission status: Runtime Error

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

###dが大きい順にソート
td.sort(key = lambda x: x[0],reverse=True)
td.sort(key = lambda x: x[1],reverse=True)

###なるべくtの種類を多いなかで、dを最大化する
###まず各tでdが最大のものを1つずつ選ぶ
eated_kind = []
eated_kind_num = 0
eated = []
for t,d in td:
  if t not in eated_kind:
    eated_kind.append(t)
    eated_kind_num += 1
    eated.append([t,d])
    
    ###tがk種類存在したら、残りのtを扱うことはない
    ###またdも最大化されているため、これ以降のtdを使うことはない
    if eated_kind_num == k:
      td = td[:i+1]
      break

not_eated = [i for i in td if i not in eated]

###tがk種類存在しなかったら、k個になるまでtに関係なくdを最大化する
if eated_kind_num < k:
  eated += not_eated[:k-eated_kind_num]
  not_eated = not_eated[k-eated_kind_num:]
  #print(eated)
  eated.sort(key = lambda x: x[1],reverse=True)
  
t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point
###種類の寿司でdの最大化をする
for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  eated = eated[:-1]
  point = max(point,point[-1]+t_point_sub+d_point_sub)
  eated_kind_num -= 1    
#print(eated)

print(point)

Submission status: Runtime Error

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

###dが大きい順にソート
td.sort(key = lambda x: x[1],reverse=True)

###なるべくtの種類を多いなかで、dを最大化する
###まず各tでdが最大のものを1つずつ選ぶ
eated_kind = []
eated = []
for i in range(n):
  t,d = td[i]
  if t not in eated_kind:
    eated_kind.append(t)
    eated.append([t,d])
    
    ###tがk種類存在したら、残りのtを扱うことはない
    ###またdも最大化されているため、これ以降のtdを使うことはない
    if len(eated_kind) == k:
      td = td[:i+1]
      break

eated_kind_num = len(eated_kind)
not_eated = [i for i in td if i not in eated]

###tがk種類存在しなかったら、k個になるまでtに関係なくdを最大化する
if eated_kind_num < k:
  eated.append(not_eated[:k-cnt])
  not_eated = not_eated[k-cnt:]
  eated.sort(key = lambda x: x[1],reverse=True)
  
###dの最大化を目指すが、tの種類が減る場合どちらのほうが良いか判断する必要がある
for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  if d < eated_last_d:
    break
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  if t_point_sub + d_point_sub > 0:
    eated = [[t,d]] + eated[:-1]
    eated_kind_num -= 1
  else:
    break
    
#print(eated)

t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point
print(point)

Submission status: Runtime Error

Solution code:
from operator import itemgetter
from collections import Counter


def solve(string):
    n, k, *td = map(int, string.split())
    dic = {i: [] for i in range(1, n + 1)}
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=itemgetter(1), reverse=True)
    for t, d in td:
        dic[t].append(d)
    base = td[:k]
    counts = Counter([_b[0] for _b in base])
    v = set(counts.keys())
    start_v = len(v)
    max_v = min(len(set(td[::2])), k)

    candidates = []
    for _td in td[k:]:
        if _td[0] not in v:
            candidates.append(_td)
            v.add(_td[0])

    ans = sum([_b[1] for _b in base]) + start_v**2
    index = -1
    tmp = ans
    for i in range(max_v - start_v):
        while counts[base[index][0]] == 1:
            index -= 1
        _t, _d = base[index]
        counts[_t] -= 1
        tmp = tmp - _d + candidates[i][1] + 2 * (start_v + i + 1) - 1
        ans = max(tmp, ans)

    return str(ans)


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Runtime Error

Solution code:
from copy import deepcopy


def solve(string):
    n, k, *td = map(int, string.split())
    dic = {i: [] for i in range(1, n + 1)}
    for t, d in zip(td[::2], td[1::2]):
        dic[t].append(d)
    dic = {_k: sorted(dic[_k], reverse=True) for _k in dic.keys()}

    def dfs(variety, b, use):
        if sum(use) + n == k:
            return len(variety)**2 + b
        tmp = []
        for _u, _k in zip(use, dic.keys()):
            if _u < len(dic[_k]) - 1:
                u = deepcopy(use)
                v = deepcopy(variety)
                u[_k - 1] += 1
                v.add(_k)
                tmp.append(dfs(v, b + dic[_k][u[_k - 1]], u))
        return max(tmp)

    return str(dfs(set([]), 0, [-1] * n))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Runtime Error

Solution code:
import heapq

INF = float('inf')
n, k = map(int, input().split())
sushi = {}

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

for s in sushi.values():
    s.sort(reverse=True)

eat_counts = {s: 0 for s in sushi.keys()}

# まだ食べてないやつ
news = []
# リピート
repeats = []
# 食べた履歴
eats = []

# 各寿司の一番おいしいやつを heap へ
for name, deli_li in sushi.items():
    # (-おいしさ, name)
    heapq.heappush(news, (-deli_li[0], name))

sum_deli = 0
variety = 0

for _ in range(k):
    # おいしいやつから食べる
    if news and (not repeats or news[0] <= repeats[0]):
        deli, name = heapq.heappop(news)
        variety += 1
    else:
        deli, name = heapq.heappop(repeats)

    sum_deli += -deli
    eat_counts[name] += 1
    heapq.heappush(eats, (-deli, name))
    if eat_counts[name] < len(sushi[name]):
        heapq.heappush(repeats, (-sushi[name][eat_counts[name]], name))

# 美味しいやつから順に食べた
# 種類増やしたらもっと得点上がるか確かめる
ans = sum_deli + variety ** 2
while True:
    if not news:
        break

    # 食べたやつポイント小さいやつから置き換えたい
    ate_min_deli, ate_min_name = heapq.heappop(eats)
    eat_counts[ate_min_name] -= 1
    if eat_counts[ate_min_name] <= 0:
        # これは種類減っちゃうからだめ
        continue

    sum_deli -= ate_min_deli

    deli, name = heapq.heappop(news)
    variety += 1
    sum_deli += -deli
    eat_counts[name] += 1

    ans = max(ans, sum_deli + variety ** 2)

print(ans)

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys
import heapq


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    heaps = {}
    for i in range(N):
        if t[i] not in heaps:
            heaps[t[i]] = []
            heapq.heapify(heaps[t[i]])
        heap = heaps[t[i]]
        heapq.heappush(heap, (-d[i], d[i]))
    used = set([])
    current = 0
    for i in range(K):
        maxIndex = -1
        max = -1
        for d in heaps:
            heap = heaps[d]
            if len(heap) == 0:
                continue
            next = heap[0][1]
            if d not in used:
                next += (len(used)+1) * (len(used)+1) - len(used)*len(used)
            if max < next:
                max = next
                maxIndex = d
        if maxIndex != -1:
            used.add(maxIndex)
            current += max
            heapq.heappop(heaps[maxIndex])
            print(f"use {maxIndex}")
    print(current)
    return


# Generated by 1.1.3 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]"
    d = [int()] * (N)  # type: "List[int]"
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys, heapq


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    group = {}
    for i in range(N):
        k = t[i]
        v = d[i]
        if k in group:
            group[k].append(v)
        else:
            group[k] = [v]
    groups = list(group.values())
    for g in groups:
        g.sort(reverse=True)
    groups.sort(reverse=True, key=lambda x: x[0])
    h = []
    ret = 0
    base = 0
    #print(groups)
    for i in range(min(len(groups), K)):
        g = groups[i]
        bonus = (i + 1) ** 2
        for j, v in enumerate(g):
            if j == 0:
                if len(h) > K - (i + 1):
                    remove = heapq.heappop(h)
                else:
                    remove = 0
                base = base - remove + v
            elif len(h) < K - (i + 1):
                heapq.heappush(h, g[j])
                base += g[j]
            elif h[0] < v:
                base = base - h[0] + v
                heapq.heapreplace(h, v)
            ret = max(ret, base + bonus)
            #print(h)
            #print("bonus : %d, base : %d, ret : %d" % (bonus, base, ret))
            #print(' --- ')
    print(ret)
    return


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()