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 operator import itemgetter
input = sys.stdin.buffer.readline

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

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

ate = {t:False for t, _ in tds}
variety = 0
deliciousness = 0
will_be_replaced = []
for t, d in tds[:K]:
    if ate[t] == False:
        ate[t] = True
        variety += 1
    else:
        will_be_replaced.append((t, d))
    deliciousness += d

max_point = deliciousness + variety ** 2

i = K
while will_be_replaced:
    t, d = will_be_replaced.pop()
    while i < N and ate[tds[i][0]]:
        i += 1
    deliciousness -= d
    deliciousness += tds[i][1]
    variety += 1
    point = deliciousness + variety ** 2
    max_point = max(max_point, point)

print(max_point)

Submission status: Runtime Error

Solution code:
import sys

readline = sys.stdin.readline
MOD = 10 ** 9 + 7
INF = float('INF')
sys.setrecursionlimit(10 ** 5)

import heapq


class PriorityQueue:
    class Reverse:
        def __init__(self, val):
            self.val = val

        def __lt__(self, other):
            return self.val > other.val

        def __repr__(self):
            return repr(self.val)

    def __init__(self, x=None, desc=False):
        if not x:
            x = []
        if desc:
            for i in range(len(x)):
                x[i] = self.Reverse(x[i])
        self._desc = desc
        self._container = x
        heapq.heapify(self._container)

    @property
    def is_empty(self):
        return not self._container

    def pop(self):
        if self._desc:
            return heapq.heappop(self._container).val
        else:
            return heapq.heappop(self._container)

    def push(self, item):
        if self._desc:
            heapq.heappush(self._container, self.Reverse(item))
        else:
            heapq.heappush(self._container, item)

    def top(self):
        if self._desc:
            return self._container[0].val
        else:
            return self._container[0]

    def sum(self):
        return sum(self._container)

    def __len__(self):
        return len(self._container)


def main():
    from operator import itemgetter

    n, k = list(map(int, readline().split()))
    sushis_original = [list(map(int, readline().split())) for _ in range(n)]

    sushis_original.sort(key=itemgetter(1))
    sushis_original.sort(key=itemgetter(0))

    new_type = 0
    prev = -1
    for i in range(n):
        cur = sushis_original[i][0]
        if prev != cur:
            new_type += 1
        if cur > new_type:
            sushis_original[i][0] = new_type
        prev = cur

    type_num = sushis_original[-1][0]

    sushis = {i: [] for i in range(1, type_num + 1)}

    for sushi_type, val in sushis_original:
        sushis[sushi_type].append(val)

    eat_sushis = PriorityQueue()
    rem_sushis = PriorityQueue(desc=True)
    rem = k

    if rem >= type_num:
        for i in range(1, type_num + 1):
            eat_sushis.push((sushis[i].pop(), i))
        rem -= type_num

        for sushi_type, vals in sushis.items():
            for val in vals:
                rem_sushis.push((val, sushi_type))
        for _ in range(rem):
            eat_sushis.push(rem_sushis.pop())
    else:
        for i in range(1, type_num + 1):
            eat_sushis.push((sushis[i].pop(), i))
        if len(eat_sushis) > k:
            discard_num = len(eat_sushi) - k
            for _ in range(discard_num):
                eat_sushis.pop()

    cur_type = min(k, type_num)
    sub_next = 2 * cur_type - 1

    while rem_sushis:
        cur_val, _ = eat_sushis.top()
        new_val, _ = rem_sushis.top()
        diff = new_val - cur_val
        if diff >= sub_next:
            eat_sushis.pop()
            eat_sushis.push(rem_sushis.pop())
            cur_type -= 1
            sub_next = 2 * cur_type - 1
        else:
            break

    ans = cur_type ** 2

    while eat_sushis:
        val, _ = eat_sushis.pop()
        ans += val

    print(ans)


if __name__ == '__main__':
    main()

Submission status: Runtime Error

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

L = [[] for _ in range(N)]
All = []
s = set([])
for i in range(N):
  t,d = map(int,input().split())
  t-=1 #0スタートに
  heapq.heappush(L[t], -d) #最大から取り出したいので負にする。
  heapq.heappush(All,[-d,t]) #すべてから大きい順に取り出せるように 
  s.add(t)

#print(L,s)
point = 0
ns = 0 #現在の種類
later = []
now = []
syurui = set([])
while len(now) < K and All:
  temp = heapq.heappop(All)
  if temp[1] not in syurui:
    ns += 1
    point -= temp[0]
    syurui.add(temp[1])
    heapq.heappush(now, -1*temp[0])
  else:
    later.append(temp)
for i in later:
  heapq.heappush(All,i)

point += ns*ns
while len(now) < K:
  temp = heapq.heappop(All)
  point -= temp[0]
  heapq.heappush(now, -1*temp[0])

#print(ns,point)
#print(now,All)
while now[0] < -1*All[0][0]:
  temp = heapq.heappop(All)
  dele = heapq.heappop(now)
  #print(temp,dele)
  upoint = point - temp[0] - dele - ns*ns + (ns-1)**2
  #print(upoint,point)
  if upoint >= point:
    point = upoint
    ns -= 1
  else:
    break
    
print(point)
#print(All)
#print(now)

Submission status: Runtime Error

Solution code:
from collections import Counter

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

TD.sort(reverse=True, key=lambda a: a[1])

A = TD[:K]

B = []
V = set()
for t, d in TD:
    if t in V:
        continue
    B.append((t, d))
    V.add(t)

cntT = Counter([t for t, _ in A])
V = set(cntT.keys())
S = sum([d for _, d in A])
B = [(t, d) for t, d in B if not t in V]

C = []
for t, d in A[::-1]:
    if cntT[t] > 1:
        C.append((t, d))
        cntT[t] -= 1

C = C[::-1]
ans = S + len(V)**2
for t, d in B:
    _, less = C.pop()
    S -= less
    S += d
    V.add(t)
    ans = max(ans, S + len(V)**2)
print(ans)

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
from operator import itemgetter
import heapq
n, k = map(int, input().split())
TD = sorted([list(map(int, input().split())) for _ in range(n)], reverse=True, key=itemgetter(1))
L = [[] for _ in range(n+1)]
P = []
a = 0
cnt = 0
F = [-float("inf")]*(n+1)
for t, d in TD[:k]:
  L[t].append(d)
  if len(L[t]) == 1:
    a += 1
  cnt += d
F[a] = cnt
for i in range(1, n+1):
  while len(L[i]) > 1:
    heapq.heappush(P, L[i].pop())
now = k
for i in range(a, n):
  while now < n and L[TD[now][0]]:
    now += 1
  if now == n:
    break
  t, d = TD[now]
  L[t].append(d)
  F[i+1] = F[i] - heapq.heappop(P) + d
ans = 0
for i in range(a, n+1):
  ans = max(ans, F[i] + i*i)
print(ans)

Submission status: Runtime Error

Solution code:
n, m = map(int, input().split())
a = list(map(int , input().split()))
t = [2,5,5,4,5,6,3,7,6]
cost = {ai: t[ai-1] for ai in a}

dp = [0] + [-1] * n

for i in range(1, n+1):
    dpi = -1
    for ai, ci in cost.items():
        if i - ci >= 0 and dp[i-ci] >= 0:
            dpi = max(dp[i-ci] * 10 + ai, dpi)
    dp[i] = dpi
print(dp[n])

Submission status: Runtime Error

Solution code:

n, k = list(map(int, input().split()))
sushi = []
for i in range(n):
    sushi.append(list(map(int, input().split())))
point = 0
sushi_kind = set()
duplication = []
sushi.sort(key=lambda x: x[1],reverse = True)
for i in range(k):
    point += sushi[i][1]
    if sushi[i][0] in sushi_kind:
        duplication.append(sushi[i])
    else:
        sushi_kind.add(sushi[i][0])
ans = point + pow(len(sushi_kind),2)
dup_index = len(duplication) - 1
for i in range(k,n):
    if sushi[i][0] in sushi_kind:
        continue
    else:
        point += sushi[i][1]
        sushi_kind.add(sushi[i][0])
        point -= duplication[dup_index][1]
    ans = max(ans, point + pow(len(sushi_kind),2))
print(ans)

Submission status: Runtime Error

Solution code:
from heapq import heappush, heappop
from collections import defaultdict

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

DT = [[d, t] for t, d  in TD]
DT.sort(reverse = True)
#print(DT)

h_in = []
cnt = defaultdict(int)

tmp = []
kiso = 0
for d, t in DT:
    if cnt[t] == 0:
        heappush(h_in, [d, t])
        kiso += d
    else:
        tmp.append([d, t])
    cnt[t] += 1

tmp = tmp[::-1]
len_h_in = len(h_in)
if len_h_in < K:
    d = K - len_h_in
    for i in range(d):
        heappush(h_in, tmp[i])
        kiso += tmp[i][0]

kind = 0
for k, v in list(cnt.items()):
    if v:
        kind += 1
#print(kiso, kind)
#print(h_in)
#print(tmp)

answer = kiso + kind ** 2
tmp = tmp[::-1]
while tmp:
    out_d, out_t = tmp.pop()
    in_d, in_t = heappop(h_in)
    
    kiso += out_d - in_d
    cnt[in_t] -= 1
    if cnt[in_t] == 0: kind -= 1
    cnt[out_t] += 1
    if cnt[out_t] == 1: kind += 1
    
    answer = max(answer, kiso + kind ** 2)

print(answer)

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
#%% for atcoder uniittest use
import sys
input= lambda: sys.stdin.readline().rstrip()
def pin(type=int):return map(type,input().split())
def tupin(t=int):return tuple(pin(t))
#%%code
import time
import random
def resolve():
    start_time=time.time()
    N,K=pin()
    neta=[]
    for i in range(N):
        x,y=pin()
        neta.append([x,y,0])
    neta.sort(key=lambda x:x[1])
    #print(neta)

    

    start_point=0
    netaset=set()
    for i in range(K):
        x,y,z=neta[N-i-1]
        netaset.add(x)
        start_point+=y
    start_point+=len(netaset)**2
    #print(start_point)


    neta.sort(key=lambda x:x[0])
    #print(neta)
    netaset=set()
    t=0
    neopoint=0
    neta2=neta[:]
    for i in range(0,N-1):
        if neta2[i+1][0]!=neta2[i][0]:
            t+=1
            s,f,q=neta[i]
            netaset.add(s)
            neopoint+=f
            neta[i][2]=1
    #print(neta)
    s,f,q=neta[N-1]
    netaset.add(s)
    neopoint+=f
    neta[i][2]=1
    
    #print(neopoint)
    a=K-t-1
    neta.sort(key=lambda x:x[1])
    #print(a,neta)
    if a>0:
        while(a>0):
            pp,qq,rr=neta.pop()
            if rr==0:
                a-=1
                netaset.add(pp)
                neopoint+=qq
    neopoint+=len(netaset)**2
    #print(max(neopoint,start_point))
    pseodomax=0
    while(time.time()-start_time<1.80):
        pseodopoint=0
        pseodoset=set()
        pseodosetlen=0
        netaset=set()

        if pseodosetlen==K:
            pseodomax=max(pseodomax,pseodopoint)
            netaset=set()
            pseodosetlen=0
            pseodoset=set()
        t=random.randrange(0,N)
        if t in pseodoset:continue
        pseodoset.add(t)
        pseodosetlen+=1
        s,f,q=neta[t]
        netaset.add(s)
        pseodopoint+=f
        
            
    print(max(start_point,neopoint,pseodomax))
#%%submit!
resolve()

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
#%% for atcoder uniittest use
import sys
input= lambda: sys.stdin.readline().rstrip()
def pin(type=int):return map(type,input().split())
def tupin(t=int):return tuple(pin(t))
#%%code

def resolve():
    N,K=pin()
    neta=[[*pin(),0] for i in range(N)]
    neta.sort(key=lambda x:x[1])
    #print(neta)

    start_point=0
    netaset=set()
    for i in range(K):
        x,y,z=neta[N-i-1]
        netaset.add(x)
        start_point+=y
    start_point+=len(netaset)**2
    #print(start_point)


    neta.sort(key=lambda x:x[0])
    #print(neta)
    netaset=set()
    t=0
    neopoint=0
    neta2=neta[:]
    for i in range(0,N-1):
        if neta2[i+1][0]!=neta2[i][0]:
            t+=1
            s,f,q=neta[i]
            netaset.add(s)
            neopoint+=f
            neta[i][2]=1
    #print(neta)
    s,f,q=neta[N-1]
    netaset.add(s)
    neopoint+=f
    neta[i][2]=1
    
    #print(neopoint)
    a=K-t-1
    neta.sort(key=lambda x:x[1])
    #print(a,neta)
    if a>0:
        while(a>0):
            pp,qq,rr=neta.pop()
            if rr==0:
                a-=1
                netaset.add(pp)
                neopoint+=qq
    neopoint+=len(netaset)**2
    #print(max(neopoint,start_point))
    
        
            
    print(max(start_point,neopoint))
#%%submit!
resolve()

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
#%% for atcoder uniittest use
import sys
input= lambda: sys.stdin.readline().rstrip()
def pin(type=int):return map(type,input().split())
def tupin(t=int):return tuple(pin(t))
#%%code

def resolve():
    N,K=pin()
    neta=[[*pin(),0] for i in range(N)]
    neta.sort(key=lambda x:x[1])
    #print(neta)

    start_point=0
    netaset=set()
    for i in range(K):
        x,y,z=neta[N-i-1]
        netaset.add(x)
        start_point+=y
    start_point+=len(netaset)**2
    #print(start_point)


    neta.sort(key=lambda x:x[0])
    #print(neta)
    netaset=set()
    t=0
    neopoint=0
    neta2=neta[:]
    for i in range(0,N-1):
        if neta2[i+1][0]!=neta2[i][0]:
            t+=1
            s,f,q=neta[i]
            netaset.add(s)
            neopoint+=f
            neta[i][2]=1
    #print(neta)
    s,f,q=neta[N-1]
    netaset.add(s)
    neopoint+=f
    neta[i][2]=1
    
    #print(neopoint)
    a=K-t-1
    neta.sort(key=lambda x:x[1])
    #print(a,neta)
    if a>0:
        while(a>0):
            pp,qq,rr=neta.pop()
            if rr==0:
                a-=1
                netaset.add(pp)
                neopoint+=qq
    neopoint+=len(netaset)**2
    #print(max(neopoint,start_point))
    
        
            
    print(max(start_point,neopoint))
#%%submit!
resolve()

Submission status: Runtime Error

Solution code:
# D - Various Sushi
import sys
from collections import defaultdict
import heapq
readline = sys.stdin.readline

dd = defaultdict(list)
N, K = map(int, readline().split())
for i in range(N):
    t, d = map(int, readline().split())
    heapq.heappush(dd[t], d*-1)

dd = list(dd.values())
dd.sort()

ans = 0
fix = 0
var = []
for i in range(len(dd)):
    fix += heapq.heappop(dd[i]) * -1
    while len(dd[i]) > 0:
        tmp = heapq.heappop(dd[i]) * -1
        heapq.heappush(var, tmp)
    while len(var) > K-i-1:
        heapq.heappop(var)
    if len(var) == K-i-1:
        ans = max(ans , fix+sum(var) + (i+1)**2)

print(ans)

Submission status: Runtime Error

Solution code:
# D - Various Sushi

from collections import defaultdict
import heapq

dd = defaultdict(list)
N, K = map(int, input().split())
for i in range(N):
    t, d = map(int, input().split())
    heapq.heappush(dd[t], d*-1)

dd = list(dd.values())
dd.sort()

ans = 0
fix = 0
var = []
for i in range(len(dd)):
    fix += heapq.heappop(dd[i]) * -1
    while len(dd[i]) > 0:
        tmp = heapq.heappop(dd[i]) * -1
        heapq.heappush(var, tmp)
    while len(var) > K-i-1:
        heapq.heappop(var)
    if len(var) == K-i-1:
        ans = max(ans , fix+sum(var) + (i+1)**2)

print(ans)

Submission status: Runtime Error

Solution code:
from heapq import heappush,heappop,heapify

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

# とりあえずネタは無視しておいしさが大きい寿司からk個を選び(hq)、満足ポイントを計算する。
# それに選ばれなかった寿司から、選ばれた寿司のネタにはないネタで、かつ、
# そのネタの中でおいしさが最大のものを選んでおく(rhq)。
# hqのうち、hqのネタ種類数が減らなように選んだおいしさが最小のものと、
# rhqのうちおいしさが最大のものを入れ替えて、満足ポイントを計算する。
# (こうすると、おいしさの合計は減るがネタの種類数は増える。)
# これを繰り返して、最大の満足ポイントを出力する。

tf = [False for _ in range(n)]
s = []
for i in range(n):
  t,d = map(int,input().split())
  t -= 1
  if tf[t]:
    s.append((0,d))
  else:
    s.append((1,d))
    tf[t] = True

s.sort(key = lambda x:-x[1])
hq = s[:k]
hq.sort()
heapify(hq)

rhq = []
for t,d in s[k:]:
  if t == 1:
    rhq.append(-d)
heapify(rhq)

ans = 0
tsum = 0
dsum = 0
for t,d in hq:
  tsum += t
  dsum += d
ans = dsum+tsum**2

while(True):
  t,d = heappop(hq)
  if t == 1:
    break
  if len(rhq) == 0:
    break
  newd = heappop(rhq)
  dsum += -newd-d
  tsum += 1
  ans = max(dsum+tsum**2,ans)
print(ans)

Submission status: Runtime Error

Solution code:
from heapq import heappush,heappop,heapify

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

# とりあえずネタは無視しておいしさが大きい寿司からk個を選び(hq)、満足ポイントを計算する。
# それに選ばれなかった寿司から、選ばれた寿司のネタにはないネタで、かつ、
# そのネタの中でおいしさが最大のものを選んでおく(rhq)。
# hqのうち、hqのネタ種類数が減らなように選んだおいしさが最小のものと、
# rhqのうちおいしさが最大のものを入れ替えて、満足ポイントを計算する。
# (こうすると、おいしさの合計は減るがネタの種類数は増える。)
# これを繰り返して、最大の満足ポイントを出力する。

tf = [False for _ in range(n)]
s = []
for i in range(n):
  t,d = map(int,input().split())
  t -= 1
  if tf[t]:
    s.append((0,d))
  else:
    s.append((1,d))
    tf[t] = True

s.sort(key = lambda x:-x[1])
hq = s[:k]
heapify(hq)

rhq = []
for t,d in s[k:]:
  if t == 1:
    rhq.append(-d)
heapify(rhq)

ans = 0
tsum = 0
dsum = 0
for t,d in hq:
  tsum += t
  dsum += d
ans = dsum+tsum**2

while(True):
  t,d = heappop(hq)
  if t == 1:
    break
  if len(rhq) == 0:
    break
  newd = heappop(rhq)
  dsum += -newd-d
  tsum += 1
  ans = max(dsum+tsum**2,ans)
print(ans)

Submission status: Runtime Error

Solution code:
from collections import defaultdict
import heapq
import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(10 ** 7)

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

DT.sort(reverse=True)
score = 0
kind = defaultdict(list)
used = [0] * (N + 1)
for d, t in DT[:K]:
    score += d
    used[t] += 1
    kind[t].append(d)

cnt = len(kind)
ans = score + (cnt * cnt)

cand = []  # 交換する可能性のあるやつ, 小さいやつから削除
for key, val in kind.items():
    val.sort(reverse=True)
    while len(val) > 1:
        heapq.heappush(cand, (val.pop(), key))


for d_in, t_in in DT[K:]:
    if not kind:
        break
    if used[t_in]:
        continue

    d_out, t_out = heapq.heappop(cand)
    used[t_in] += 1
    used[t_out] -= 1
    cnt += 1
    score = score - d_out + d_in
    ans = max(ans, score + (cnt * cnt))

print(ans)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(100000000)

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

    sushi.sort(key = lambda x:x[1],reverse = True)
    ans = 0
    dic = {}
    change = []
    for a,b in sushi[:K]:
        if a in dic:
            dic[a].append(b)
            change.append((a,b))
        else:
            dic[a] = [b]
        ans += b

    various = len(dic)
    ans += various ** 2
    ret = ans
    used_type = {a for a in dic.keys()}
    suggestion = [p for p in sushi[K:] if p[0] not in used_type]
    for a,b in suggestion:
        if a in used_type:
            continue
        dic[change[-1][0]].pop()
        ret -= change[-1][1]
        ret += b + 2*various + 1
        dic[a] = b
        change.pop()
        various += 1
        used_type.add(a)
        ans = max(ans,ret)
    print(ans)
if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline

N, K = map(int, input().split())
nums = []
for i in range(N):
    t, d = map(int, input().split())
    nums.append((d, t))
nums.sort(reverse=True)
se = set()
kind = 0
a1 = []
a2 = []
rest = []
for i, n in nums:
    if i < K:
        if not n[1] in se:
            se.add(n[1])
            kind += 1
            a1.append(n)
        else:
            a2.append(n)
    else:
        rest.append(n)
points = 0
for (d, t) in a1: points += d
for (d, t) in a2: points += d
ans = points + pow(kind, 2)
a2.sort()
for i, (d, t) in enumerate(rest):
    if t in se: continue
    if i >= len(a2): break
    points -= a2[i]
    points += d
    kind += 1
    ans = max(ans, points+pow(kind, 2))
print(ans)

Submission status: Runtime Error

Solution code:
from collections import defaultdict
from heapq import heappush, heappop
N,K,*L = map(int, open(0).read().split())
dic = defaultdict(list)
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  dic[t].append(d)
  heappush(ls,(-d,t))
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = heappop(ls)
  ans -= d
  if t not in S:
    num += 1
    S.add(t)
  else:
    heappush(pre,-d)
ans += num*num
m = ans
for i in range(N-K):
  d,t = heappop(ls)
  if t in S:
    continue
  S.add(t)
  m -= heappop(pre)
  m += -d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

Submission status: Runtime Error

Solution code:
from collections import deque
n,k=map(int,input().split())
td=[[0]*2 for _ in range(n)]
for i in range(n):
    t,d=map(int,input().split())
    td[i][0],td[i][1]=d,t
td.sort(reverse=True)
di=set()
sa=deque()
ans=0
for d,t in td[:k]:
    if t in di:
        sa.append(d)
    else:
        di.add(t)
    ans+=d

s=[0]*(k+1)
cnt=len(di)
s[cnt]=ans
for d,t in td[k:]:
    if t in di:
        continue
    else:
        x=sa.popleft()
        ans+=d-x
        di.add(t)
        cnt+=1
        s[cnt]=ans

for i in range(k+1):
    s[i]+=i*i

print(max(s))

Submission status: Runtime Error

Solution code:
import heapq
n,k=map(int,input().split())
kind=[0]*n
s=set()
dt=[]
ans=0
for _ in range(n):
    t,d=map(int,input().split())
    ans+=d
    t-=1
    dt.append((d,t))
    s.add(t)
    kind[t]+=1
x=len(s)
ans+=x**2
sing,comp=[],[]
for d,t in dt:
    if kind[t]==1:
        heapq.heappush(sing,(d,t))
    elif kind[t]>=2:
        heapq.heappush(comp,(d,t))
for _ in range(n-k):
    ds,ts=heapq.heappop(sing)
    dc,tc=heapq.heappop(comp)
    if ds+x**2-(x-1)**2>dc:
        ans-=dc
        kind[tc]-=1
        if kind[tc]>=2:
            heapq.heappush(comp,(dc,tc))
        else:
            heapq.heappush(sing,(dc,tc))
        heapq.heappush(sing,(ds,ts))
    else:
        ans-=ds+x**2-(x-1)**2
        x-=1
        kind[ts]-=1
        heapq.heappush(comp,(dc,tc))
print(ans)

Submission status: Runtime Error

Solution code:
from heapq import heappush, heappop, heapify
import sys
input = sys.stdin.buffer.readline

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

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

heapify(sushi)

fixed = {}
ate = []

for i in range(k):
    d, t = heappop(sushi)
    d = -d

    if t not in fixed:
        fixed[t] = d
    else:
        ate.append(d)

ate.sort(reverse=True)


def get_score():
    return sum(fixed.values()) + sum(ate) + len(fixed) * len(fixed)


score = get_score()

while sushi:
    d, t = heappop(sushi)
    d = -d

    if t not in fixed:
        fixed[t] = d
        del ate[-1]
        score = max(score, get_score())
        if len(fixed) == k:
            break

print(score)

Submission status: Runtime Error

Solution code:
from heapq import heappush, heappop
import sys
input = sys.stdin.buffer.readline

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

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

fixed = {}
ate = []

for i in range(k):
    d, t = heappop(sushi)
    d = -d

    if t not in fixed:
        fixed[t] = d
    else:
        heappush(ate, d)


def get_score():
    return sum(fixed.values()) + sum(ate) + len(fixed) * len(fixed)


score = get_score()

while sushi:
    d, t = heappop(sushi)
    d = -d

    if t not in fixed:
        fixed[t] = d
        heappop(ate)
        score = max(score, get_score())
        if len(fixed) == k:
            break

print(score)

Submission status: Runtime Error

Solution code:
# input()
# int(input())
# map(int, input().split())
# list(map(int, input().split()))

import numpy as np
inf = 10**9+7
mod = 10**9+7
n, k = map(int, input().split())
g_val = [list(map(int, input().split())) for i in range(n)]
netalist = [False for i in range(n+1)]
netanum = [0 for i in range(n+1)]
g_val.sort(key=lambda x: x[1])
used_list = []
ans = [0 for i in range(k+1)] # i種類食べた時の最大値
preans = 0
for j in range(k):
    if g_val:
        g,v = g_val.pop()
        netalist[g] = True
        netanum[g]+=1
        used_list.append((v,g))
        preans+=v

kind = sum(list(map(int, netalist)))
ans[kind] = preans+(kind*kind)

import heapq # 優先度付きキュー(最小値取り出し)
heapq.heapify(used_list)   # ヒープ化


while used_list:
    v,g = heapq.heappop(used_list) # 最小値の取り出し
    if netanum[g] > 1:
        while g_val:
            gnew,vnew = g_val.pop()
            if netanum[gnew] == 0:
                ans[kind+1] = ans[kind]-v-(kind*kind)+vnew+((kind+1)*(kind+1))
                kind+=1
                netanum[g] -= 1
                netanum[gnew] += 1
                break
    if not g_val:
        break
print(max(ans))

Submission status: Runtime Error

Solution code:
from collections import defaultdict
import heapq
N,K=map(int,input().split())

hq=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  heapq.heappush(hq,(-d,t))
  tset_all.add(t)
#print(heapq)

hq_K=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  md,t=heapq.heappop(hq)
  heapq.heappush(hq_K,((-md,t)))
  dsum-=md
  tdic[t]+=1  

t0=len(tdic)
max_answer=dsum+(t0**2)
#print(t0,max_answer)

for i in range(t0+1,len(tset_all)+1):  
  loop_flg=True
  while(loop_flg):
    d,t=heapq.heappop(hq_K)
    if tdic[t]>1:
      tdic[t]-=1
      dsum-=d
      
      while(True):
        md2,t2=heapq.heappop(hq)
        if tdic[t2]==0:
          dsum-=md2
          loop_flg=False
          break
          
  answer_i=dsum+(i**2)
  #print(i,answer_i)
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

Submission status: Runtime Error

Solution code:
from collections import deque
n,k = map(int,input().split())
td = [list(map(int,input().split())) for i in range(n)]
td.sort(key=lambda x:x[1],reverse=True)
a = 0
ans = []
kind = [-1]*n
c = 0
min_p = deque()
for i in range(n):
    eat = False
    if i<k:
        a += td[i][1]
        eat = True
    else:
        if min_p!=[]:
            if kind[td[i][0]-1]<0:
                a -= (min_p[-1]-td[i][1])
                a += (c+1)**2-c**2
                ans.append(a)
                min_p.popleft()
                eat = True
    if eat:
        if kind[td[i][0]-1]<0:
            kind[td[i][0]-1]=1
            c+=1
        else:min_p.append(td[i][1])
    if i==k-1:
        a += c**2
        ans.append(a)
print(max(ans))

Submission status: Runtime Error

Solution code:
n,k=map(int,input().split())
td=[]
for _ in range(n):
  td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
td_core=[]
td_sel=[]
td_sub=[]
cat=[0]*(n+1)
for i in range(n):
  if i < k and cat[td[i][0]]==0:
    td_core.append(td[i])
    cat[td[i][0]]=1
  elif i < k:
    td_sel.append(td[i])
  elif cat[td[i][0]]==0:
    td_sub.append(td[i])
    cat[td[i][0]]=1
# i >= k and cat[td[i][0]]==1 となるものはすでにtd_core,td_sel,td_sub内ものと種類も被っている上に美味しさも低い
# 後の工程では種類を増やす処理なので、同じ種類のネタは美味しさが最大のもの一つで十分。なので捨てる

# dの大きいじゅんに貪欲にネタを取得したときの満足ポイント
kiso=sum([x[1] for x in td_core])+sum(x[1] for x in td_sel)
bns=len(td_core)
ans=kiso+bns**2
# この状態からネタの種類を増やしていき最大が出てくるか検証
# td_sel内のネタは同一ネタが複数あり、そのネタ内で最大じゃないものが美味しさの大きい順に入っている。
# td_subにはtd_core,tdselにないネタの最大美味しさの寿司が入っている。
td_sel.sort(key=lambda x:x[0])
for i in range(len(td_sel)):
  kiso-=td_sel[i][1]
  kiso+=td_sub[i][1]
  bns+=1
  ans=max(ans,kiso+bns**2)
print(ans)

Submission status: Runtime Error

Solution code:
n,k=map(int,input().split())
td=[list(map(int,input().split()))[::-1]for _ in range(n)]
td.sort(reverse=1)
s=set()
ss=set()
a=[]
b=[]
c=[]
for i in range(n):
  if i<k:
    d,t=td[i]
    if t in s:
      b.append(d)
    else:
      a.append(d)
      s.add(t)
      ss.add(t)
  else:
    d,t=td[i]
    if t not in ss:
      c.append(d)
      ss.add(t)
c=c[::-1]
ls=len(s)
su=sum(a)+sum(b)
ans=ls*ls+su
for i in range(len(ss)-len(s)):
  ls+=1
  su-=b.pop()
  su+=c.pop()
  ans=max(ans,ls*ls+su)
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])
#print(td)
ans=0
kinds=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kinds:
        kinds[td[i][0]].append(td[i][1])
    else:
        kinds[td[i][0]]=[td[i][1]]
ans+=(len(kinds)**2)
#kindsの中はソートされている状況
#print(kinds)
ans_=[ans]
for i in range(k,n):
    if td[i][0] not in kinds:
        mi=[1000000000,1000000000]
        for j in kinds:
            if len(kinds[j])>1:
                if mi[1]>kinds[j][-1]:
                    mi=[j,kinds[j][-1]]
        if mi=[1000000000,1000000000]:
            break
        else:
            ans=ans-len(kinds)**2+(len(kinds)+1)**2-mi[1]+td[i][1]
            kinds[mi[0]].pop(-1)
            kinds[td[i][0]]=[td[i][1]]
            ans_.append(ans)
#print(kinds)
print(max(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])
#print(td)
ans=0
kinds=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kinds:
        kinds[td[i][0]].append(td[i][1])
    else:
        kinds[td[i][0]]=[td[i][1]]
ans+=(len(kinds)**2)
#kindsの中はソートされている状況
#print(kinds)
ans_=[ans]
for i in range(k,n):
    if td[i][0] not in kinds:
        mi=[1000000000,1000000000]
        for j in kinds:
            if len(kinds[j])>1:
                if mi[1]>kinds[j][-1]:
                    mi=[j,kinds[j][-1]]
        if mi==[-1,-1]:
            break
        else:
            ans=ans-len(kinds)**2+(len(kinds)+1)**2-mi[1]+td[i][1]
            kinds[mi[0]].pop(-1)
            kinds[td[i][0]]=[td[i][1]]
            ans_.append(ans)
#print(kinds)
print(max(ans_))

Submission status: Runtime Error

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []
for key in sushi:
    sushi[key].sort(reverse=True)
    first.append(sushi[key][0])
    second += sushi[key][1:]

first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))


ans = 0
for i in range(1, K+1):
    if K - i <= len(second_cum)-1:
        ans = max(ans, first_cum[i] + second_cum[K-i] + i*i)

print(ans)