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: Accepted

Solution code:
def getN():
    return int(input())
def getNM():
    return map(int, input().split())
def getList():
    return list(map(int, input().split()))
def getArray(intn):
    return [int(input()) for i in range(intn)]
def input():
    return sys.stdin.readline().rstrip()
def rand_N(ran1, ran2):
    return random.randint(ran1, ran2)
def rand_List(ran1, ran2, rantime):
    return [random.randint(ran1, ran2) for i in range(rantime)]
def rand_ints_nodup(ran1, ran2, rantime):
  ns = []
  while len(ns) < rantime:
    n = random.randint(ran1, ran2)
    if not n in ns:
      ns.append(n)
  return sorted(ns)

def rand_query(ran1, ran2, rantime):
  r_query = []
  while len(r_query) < rantime:
    n_q = rand_ints_nodup(ran1, ran2, 2)
    if not n_q in r_query:
      r_query.append(n_q)
  return sorted(r_query)

from collections import defaultdict, deque, Counter
from sys import exit
from decimal import *
import heapq
from math import sqrt, pi
from fractions import gcd
import random
import string
import copy
from itertools import combinations, permutations, product
from operator import mul, itemgetter
from functools import reduce
from bisect import bisect_left, bisect_right

import sys
sys.setrecursionlimit(1000000000)
mod = 10 ** 9 + 7

#############
# Main Code #
#############

N, K = getNM()
L = [getList() for i in range(N)]

# 寿司を美味しさの順に並べる
various = defaultdict(int)
L.sort(key = lambda i: i[1], reverse = True)

pool_der = []
ans = 0
# 美味しい順にK個取る
for i in range(K):
    t, d = L[i]
    # 種類をレコードする
    various[t] += 1
    # もし同じ種類のものを既に取っていたら
    if various[t] > 1:
        pool_der.append(d)
    ans += d
# 美味しい順にソートする
pool_der.sort(reverse = True)

# 現在の種類の数
now = len(various.items())
ans += now ** 2
L = L[K:]

# まだ出ていない種類の寿司の中で最大値を取るものをレコードする
pool_var = []
for i in L:
    # まだ登場していないなら
    # 2番手以下は必要ない
    if various[i[0]] == 0:
        pool_var.append(i[1])
        various[i[0]] += 1
# 後ろから美味しい順にソートする
pool_var.sort()

# 取っても種類数が変わらない寿司を取り、新たな種類の寿司を取る
point = ans
M = min(len(pool_der), len(pool_var))

for _ in range(M):
    der = pool_der.pop()
    var = pool_var.pop()

    point += (var - der)
    point += 2 * now + 1
    now += 1
    ans = max(ans, point)
print(ans)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
from operator import itemgetter as it
l=sorted([list(map(int,input().split()))for i in range(n)],reverse=1,key=it(1))
from collections import defaultdict
c=defaultdict(int)
from heapq import heapify, heappop, heappush
now=[]
S=0
kind=0
for i in range(k):
    a,s=l[i]
    S+=s
    c[a]+=1
    if c[a]>1:heappush(now,(s,a))
    kind+=c[a]==1
ans=S+kind**2
for i in range(k,n):
    if not now:break
    x,y=l[i]
    if c[x]:continue
    s,a=heappop(now)
    c[x]+=1
    kind+=1
    S+=y-s
    ans=max(ans,S+kind**2)
print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
tds = []
for i in range(n):
    t, d = map(int, input().split())
    tds.append((d,t))
tds.sort(reverse=True)
cnt = k
hand = []
hand_throw = []
hs = {}
subs = []
rs = {}
v_t = {}
res = 0
for i in range(n):
    if cnt > 0:
        if tds[i][1] in hs:
            hs[tds[i][1]] += 1
            hand_throw.append(tds[i][0])
        else:
            hs[tds[i][1]] = 1
            hand.append(tds[i][0])
        cnt -= 1
        res += tds[i][0]
    else:
        if not tds[i][1] in rs and not tds[i][1] in hs:
            rs[tds[i][1]] = 1
            subs.append(tds[i][0])
    if not tds[i][1] in v_t:
        v_t[tds[i][1]] = 1
subs.sort()
v_hs = len(hs)
v_t = len(v_t)
res += pow(v_hs,2)
final_res = res
while v_hs+1 <= k and v_hs+1 <= v_t:
    res += subs.pop() - hand_throw.pop() - pow(v_hs,2) + pow(v_hs+1,2)
    final_res = max(final_res,res)
    v_hs += 1
print(final_res)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((t-1, d))

TD.sort(key=lambda x:-x[1])

import heapq

Tcount = [0]*N
chosen = []
heapq.heapify(chosen)

tot = 0
num_kinds = 0
ans = [0]*(N+1)

for i in range(K):
    t, d = TD[i]
    tot += d
    Tcount[t] += 1
    if Tcount[t] == 1: 
        num_kinds += 1
    if Tcount[t] > 1:
        heapq.heappush(chosen, d)

ans[num_kinds] = tot + num_kinds**2

for i in range(K, N):
    if len(chosen) == 0: break
    t, d = TD[i]
    if Tcount[t] > 0: continue
    Tcount[t] += 1
    num_kinds += 1
    tot = tot - heapq.heappop(chosen) + d
    ans[num_kinds] = tot + num_kinds**2
    
print(max(ans))

Submission status: Accepted

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from operator import itemgetter
from collections import defaultdict,deque

n,k = map(int,readline().split())
chk = []
for i in range(n):
  t,d = map(int,readline().split())
  chk.append((t,d))
chk = sorted(chk,key=itemgetter(1),reverse=True)

T = defaultdict(int)
d = deque()
s = set()
ans = []
aa = 0
for i in range(k):
  p,q = chk[i]
  T[p] += 1
  if T[p] > 1:
    d.append(q)
  s.add(p)
  aa += q
  
l = len(s)
aa += l**2
if k == n or l == k:
  print(aa)
  exit()
  
ans = aa
cnt = k-l
for p,q in chk[k:]:
  if T[p] == 0:
    T[p] += 1
    P = d.pop()
    aa -= P
    aa += q+2*l+1
    l += 1
    cnt -= 1
    ans = max(ans,aa)
  if cnt == 0:
    break
    
print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
from collections import defaultdict

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

data.sort(key=lambda x:-x[1])
KISO = 0
EAT = defaultdict(int)
for i in range(K):
  t,v = data[i]
  EAT[t] += 1
  KISO += v
SHURUI = len(EAT)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if EAT[a] == 1:
    l -= 1
    continue
  if EAT[c] != 0:
    r += 1
    continue
  EAT[a] -= 1
  EAT[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  l -= 1
  r += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Accepted

Solution code:

# -*-coding: utf-8 -*-

import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [0], [0]
    st = set()

    for t, d in S:
        if t in st:
            #累積報酬和のリスト
            y_0.append(y_0[-1]+d)
        else:
            st.add(t)
            y_1.append(y_1[-1]+d)
    

    max_sum = -1

    for i in range(1,k+1):
        if k-i >= len(y_0):
            continue
        if i >= len(y_1):
            break
        max_sum = max(max_sum, y_0[k-i]+y_1[i]+i**2)
    
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
S = [tuple(map(int, input().split())) for _ in range(N)]

S.sort(key=lambda x: -x[1])
st = set()
first_acc, other_acc = [0], [0]
for t, d in S:
    if t in st:
        other_acc.append(other_acc[-1] + d)
    else:
        st.add(t)
        first_acc.append(first_acc[-1] + d)

ans = 0
for i in range(K + 1):
    if len(first_acc) <= i or len(other_acc) <= K - i:
        continue
    ans = max(ans, first_acc[i] + other_acc[K - i] + i * i)
print(ans)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
TD = [tuple(map(int,input().split())) for i in range(N)]
TD.sort(key=lambda x:x[1])

import heapq
hq = []
heapq.heapify(hq)
selected = set()
ans = 0
for _ in range(K):
    t,d = TD.pop()
    if t in selected:
        heapq.heappush(hq, d)
    else:
        selected.add(t)
    ans += d
ans += len(selected)**2

tmp = ans
while TD:
    t,d = TD.pop()
    if t in selected: continue
    selected.add(t)
    if len(hq)==0: break
    tmp -= heapq.heappop(hq)
    tmp += d
    tmp += len(selected)*2 - 1
    ans = max(ans ,tmp)
print(ans)

Submission status: Accepted

Solution code:
from heapq import *

n, k = map(int,input().split())
T = [[] for i in range(n)]
for i in range(n):
    t,d = map(int,input().split())
    T[t-1].append(d)
Top = []
for i in range(n):
    if T[i]:
        T[i].sort(reverse=True)
        Top.append((T[i][0],i))
Top.sort(reverse = True)

various = len(Top)


ans = 0
s = 0
sushi = []
heapify(sushi)
l = 0

for j,(t,i) in enumerate(Top[:k]):
    s += t
    if l + j + 1 > k:
        s -= heappop(sushi)
        l -= 1
    for new in T[i][1:]:
        if l + j + 1 >= k:
            s += new
            s -= heappushpop(sushi,new)
        else:
            s += new
            l += 1
            heappush(sushi, new)
    ans = max(ans, s + (j+1)**2)
print(ans)

Submission status: Accepted

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

I = []
for _ in range(N):
  t,d = map(int, input().split())
  I.append((d,t))
  
I.sort(reverse=True)

# heap_can_delete
h_d = []
# heap_can_add
h_a = []

wk = 0
kind = 0
# 出現したことのあるネタ
ap = set()
for i in range(K):
  d,t = I[i]
  wk += d
  if t not in ap:
    ap.add(t)
    kind += 1
  else:
    heapq.heappush(h_d, (d,t))

for i in range(K,N):
  d,t = I[i]
  if t not in ap:
    ap.add(t)
    # maxをpopさせるために、価値を-1かけておく
    heapq.heappush(h_a, ((-1) * d, t))
  
ans = wk + kind ** 2
while len(h_d) > 0 and len(h_a) > 0:
  d1,t1 = heapq.heappop(h_d)
  d2,t2 = heapq.heappop(h_a)
  d2 = (-1) * d2
  wk = wk - d1 + d2
  kind += 1
  ans = max(ans, wk + kind ** 2)
  
print(ans)

Submission status: Accepted

Solution code:
import heapq

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

eaten = [False]*N

#d1: 種類が増えないもの同士、d2: 同じ種類の中で最大の組み合わせ
d1, d2 = [0], [0]

for i in range(N):
    t, d = sushi[i]
    t -= 1
    if eaten[t]:
        d1.append(d1[-1]+d)
    else:
        eaten[t] = True
        d2.append(d2[-1]+d)


ans = 0
for i in range(1, K+1):
    if K-i >= len(d1):
        continue
    if i >= len(d2):
        continue
    ans = max(ans, d1[K-i]+d2[i]+i**2)
print(ans)

Submission status: Accepted

Solution code:
import sys
def input():return sys.stdin.readline()[:-1]
def N(): return int(input())
def NM():return map(int,input().split())
def L():return list(NM())
def LN(n):return [N() for i in range(n)]
def LL(n):return [L() for i in range(n)]
n,k=NM()
l=LL(n)
l.sort(key=lambda x:-x[1])
from heapq import heappush,heappop,heapify
s=set()
q=[]
oisisa=0
syu=0
for t,d in l[:k]:
    if t in s:
        heappush(q,d)
    else:
        s.add(t)
        syu+=1
    oisisa+=d
ans=oisisa+syu**2
for t,d in l[k:]:
    if q and not t in s:
        syu+=1
        oisisa+=d-heappop(q)
        s.add(t)
    ans=max(ans,oisisa+syu**2)
print(ans)

Submission status: Accepted

Solution code:
from heapq import *
import sys
sys.setrecursionlimit(10**6)
I=sys.stdin.readlines()
N,K=map(int,I[0].split())
t,d=0,0
X=[]
for i in range(N):
  t,d=map(int,I[i+1].split())
  X.append((d,t))
X.sort()
Q=[]
heapify(Q)
C=[0]*N
V=0
S=0
for i in range(K):
  heappush(Q,X[-1])
  if C[X[-1][1]-1]==0:
    V+=1
  S+=X[-1][0]
  C[X[-1][1]-1]+=1
  del X[-1]
def f(q,c,s,x,v):
  y=0
  r=s+v*v
  for i in range(len(q)):
    y=heappop(q)
    if c[y[1]-1]>=2:
      s-=y[0]
      c[y[1]-1]-=1
      break
    if len(q)==0:
      return r
  for i in range(len(x)):
    if c[x[-1][1]-1]==0:
      heappush(q,x[-1])
      s+=x[-1][0]
      c[x[-1][1]-1]=1
      v+=1
      del x[-1]
      return max(r,f(q,c,s,x,v))
    del x[-1]
  return r
print(f(Q,C,S,X,V))

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mii = []
	sushi_kii = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mii.append(val)
				else:
					sushi_kii.append(val)

	sushi_mii.sort()
	sushi_mii = sushi_mii[::-1]
	sushi_kii.sort()
	sushi_kii = sushi_kii[::-1]

	sushi_mi = []
	sushi_ki = []
	for i,val in enumerate(sushi_mii):
		if i == 0:
			sushi_mi.append(val)
		else:
			tmp = sushi_mi[i-1]
			sushi_mi.append(tmp + val)

	for i,val in enumerate(sushi_kii):
		if i == 0:
			sushi_ki.append(val)
		else:
			tmp = sushi_ki[i-1]
			sushi_ki.append(tmp + val)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		tmp += sushi_mi[i-1]
		x = i

		if K - i > 0:
			tmp += sushi_ki[K - i-1]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Submission status: Accepted

Solution code:
import collections
N,K = map(int,input().split())
dic = {}
A = []
for _ in range(N):
    t,d = map(int,input().split())
    A.append([t,d])
A.sort(key = lambda x:x[1], reverse = True)
val = 0
set_ = set()
duplicated = collections.deque([])
for t,d in A[:K]:
    val += d
    if t in set_:
        duplicated.append(d)
    else:
        set_.add(t)
now = val + len(set_)**2

nx = collections.deque(A[K:])
ans = now
c = len(set_)
while len(set_) < K and nx:
    i,j = nx.popleft()
    if i in set_:
        continue
    d = duplicated.pop()
    diff = j-d + 2*c+1
    set_.add(i)
    now += diff
    ans =max(ans,now)
    c += 1
print(ans)

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
sushi = [[int(i) for i in input().split()] for _ in range(n)]
sushi.sort(key=lambda x: x[1],reverse=True)
eaten = [False]*n
d1 = [0]
d2 = [0]
for i in range(n):
    t,d = sushi[i]
    t -= 1
    if eaten[t]: d1.append(d1[-1]+d)
    else:
        eaten[t] = True
        d2.append(d2[-1]+d)
# print(d1)
# print(d2)
ans = 0
for i in range(1,k+1):
    if k-i >= len(d1): continue
    if i >= len(d2): continue
    ans = max(ans, d1[k-i] + d2[i] + i*i)
print(ans)

Submission status: Accepted

Solution code:
import heapq
from collections import Counter
inf = 10**10
n, k = map(int, input().split())

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

for dd in dset:
    if dd: dd.sort(reverse=True)
    else: dd.append(-inf)

dset.sort(key= lambda x: -x[0])

cur = 0
hq = []
# k種類の先頭を除いて格納
for i in range(k):
    cur+=dset[i][0]
    for d in dset[i][1:]: heapq.heappush(hq, -d)

# 残りの種類のすべてを格納
for dd in dset[k:]:
    for d in dd:
        heapq.heappush(hq, -d)

ans = cur + k*k
for i in range(k-1, -1, -1):
    v = dset[i][0] # 削除候補
    if hq:
        w = -heapq.heappop(hq)
    else:
        break
    if v < w:
        cur += w - v
        # heapq.heappush(hq, -v)
    else:
        heapq.heappush(hq, -w)
    ans = max(ans, cur + i*i)
print(ans)

Submission status: Accepted

Solution code:
#from collections import deque,defaultdict
printn = lambda x: print(x,end='')
inn = lambda : int(input())
inl   = lambda: list(map(int, input().split()))
inm   = lambda:      map(int, input().split())
ins = lambda : input().strip()
DBG = True # and False
BIG = 10**18
R = 10**9 + 7

def ddprint(x):
  if DBG:
    print(x)

n,k = inm()
dt = []
for i in range(n):
    t,d = inm()
    dt.append((d,t))
dt.sort(reverse=True)
a = dt[:k] # oisisa max, ignore variety

vals = [z[0] for z in a]
dsum = sum(vals)
typs = set([z[1] for z in a])
ntyp = len(typs)
mx = dsum+ntyp**2 # satisfy for a

b = dt[k:]
h = {}
for d,t in b:
    if t in typs:
        continue
    elif t not in h:
        h[t] = d
    else:
        h[t] = max(h[t],d)
c = [h[t] for t in h]
c.sort(reverse=True)
# for types not in a, max oisisa

h = {}
for d,t in a:
    if t not in h:
        h[t] = d
    else:
        h[t] = max(h[t],d)
vals = []
for d,t in a:
    if t in h and d==h[t]:
        del h[t]
    else:
        vals.append(d)
vals.sort()
# oisisa list for second or below for each type

for i in range(min(len(vals),len(c))):
    dsum += c[i]-vals[i]
    ntyp += 1
    mx = max(mx,dsum+ntyp**2)
print(mx)

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
t=[list(map(int,input().split())) for _ in range(N)]
t.sort(key=lambda x:-x[1])
ans=0
dp=set()
p=[]
for i in range(K):
  ans+=t[i][1]
  if t[i][0] in dp:
    p.append(t[i][1])
  dp.add(t[i][0])
ans+=len(dp)**2
wer=ans
while K<N and p:
  if t[K][0] not in dp:
    wer=wer+t[K][1]-p.pop()-len(dp)**2+(len(dp)+1)**2
    ans=max(ans,wer)
    dp.add(t[K][0])
  K+=1
print(ans)

Submission status: Accepted

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

n, k = map(int, input().split())
s = []
for _ in range(n):
    t, d = map(int, input().split())
    s.append((d, t))
s.sort()
bp = 0
x = 0
c = Counter()
cand = []
for _ in range(k):
    d, t = s.pop()
    bp += d
    if c[t] == 0:
        x += 1
    c[t] += 1
    heappush(cand, (d, t))
ans = bp + x * x
while s:
    d, t = s.pop()
    if c[t] == 0:
        while cand:
            dd, td = heappop(cand)
            if c[td] == 1:
                continue
            else:
                c[td] -= 1
                bp -= dd
                break
        else:
            break
        bp += d
        x += 1
        c[t] = 1
        ans = max(ans, bp + x * x)
print(ans)

Submission status: Accepted

Solution code:
import math
#import numpy as np
import queue
from collections import deque,defaultdict
import heapq as hpq
from sys import stdin,setrecursionlimit
#from scipy.sparse.csgraph import dijkstra
#from scipy.sparse import csr_matrix
ipt = stdin.readline
setrecursionlimit(10**7)

def main():
    n,k = map(int,ipt().split())
    mat = [-1]*n
    susis = []
    smo = 0
    knds = 0
    for _ in range(n):
        t,d = map(int,ipt().split())
        t -= 1
        if mat[t] < d:
            mat[t] = d
        susis.append((d,t))
    susis.sort(reverse=True)
    kh = []
    res = []
    for i in range(k):
        di,ti = susis[i]
        smo += di
        if di == mat[ti]:
            mat[ti] = -1
            knds += 1
        else:
            hpq.heappush(res,di)
    for i in mat:
        if i != -1:
            hpq.heappush(kh,-i)

    ma = smo + knds**2
    while len(res) and len(kh):
        knds += 1
        smo -= hpq.heappop(res)
        smo -= hpq.heappop(kh)
        if ma < smo+knds**2:
            ma = smo+knds**2

    print(ma)


    return

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
# various sushi
import heapq
N, K = map(int, input().split())
que = []
for i in range(N):
    t, d = map(int, input().split())
    que.append((t, d))

que.sort(key=lambda x: -x[1])
NOW = que[:K]

N_set = set()

sums = 0

cnt = [0 for i in range(N+1)]  # ねたxに何個いるか
ng,ok,sub = [], [], []
heapq.heapify(ng)
heapq.heapify(ok)
heapq.heapify(sub)

for index, value in NOW:
    sums += value
    cnt[index] += 1
    heapq.heappush(ng, (value, index))
    N_set.add(index)
while ng:
    value, index = heapq.heappop(ng)
    if cnt[index] > 1:
        heapq.heappush(ok, (value, index))
        cnt[index] -= 1
    elif cnt[index] == 1:
        heapq.heappush(sub, (value, index))

ng = sub

# ng とっちゃダメなやつ
# ok とってもいいやつ

S = len(N_set)
ans = S**2+sums


# これより更新することはあるか？

for i in range(K, N):
    taste, value = que[i]
    if taste in N_set:
        pass
    else:
        #taste not in N_set
        if ok:
            min_value, m_taste = heapq.heappop(ok)
            sums -= min_value
            sums += value
            S += 1
            cnt[taste] += 1
            N_set.add(taste)
            cnt[m_taste] -= 1
            heapq.heappush(ng, (value, taste))
            ans = max(ans, S**2+sums)

print(ans)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3

from collections import deque
n, k = list(map(int, input().split()))

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

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

set_ = set()
val = 0
duplicated = deque([])

for t_tmp, d_tmp in td[:k]:
    val += d_tmp
    if t_tmp in set_:
        duplicated.append(d_tmp)
    else:
        set_.add(t_tmp)

now = val + len(set_)**2


nx = deque(td[k:])
ans = now
c = len(set_)
while len(set_) < k and nx:
    i, j = nx.popleft()
    if i in set_:
        continue
    d = duplicated.pop()
    diff = j-d + 2 * c + 1
    set_.add(i)
    ans = max(ans, now + diff)
    now += diff
    c += 1
print(ans)

Submission status: Accepted

Solution code:
def resolve():
    N, K = list(map(int, input().split()))
    TD = sorted([list(map(int, input().split())) for _ in range(N)], reverse=True, key=lambda x: x[1])
    eat_types = set()
    eat_dups = []
    total = 0
    for i in range(K):
        typ, point = TD[i]
        total += point
        if typ in eat_types:
            eat_dups.append(TD[i])
        else:
            eat_types.add(typ)
    ans = total + len(eat_types)**2

    didx = len(eat_dups)-1
    for i in range(K, N):
        typ, point = TD[i]
        if typ in eat_types:
            continue
        if didx < 0:
            break
        total -= eat_dups[didx][1]
        total += point
        didx -= 1
        eat_types.add(typ)
        ans = max(ans, total+len(eat_types)**2)
    print(ans)
    
    
    
    

if '__main__' == __name__:
    resolve()

Submission status: Accepted

Solution code:
import sys
from collections import defaultdict
from heapq import heappush, heappop
from operator import itemgetter

sys.setrecursionlimit(10 ** 7)
rl = sys.stdin.readline


def solve():
    N, K = map(int, rl().split())
    td = [list(map(int, rl().split())) for _ in range(N)]
    
    td.sort(key=itemgetter(1), reverse=True)
    hq = []
    counter = defaultdict(int)
    t_set = set()
    d_sum = 0
    for i in range(K):
        ti, di = td[i]
        t_set.add(ti)
        counter[ti] += 1
        if 1 < counter[ti]:
            heappush(hq, di)
        d_sum += di
    
    ans = d_sum + len(t_set) * len(t_set)
    for i in range(K, N):
        ti, di = td[i]
        if ti in t_set:
            continue
        if not hq:
            break
        dj = heappop(hq)
        d_sum -= dj
        d_sum += di
        t_set.add(ti)
        ans = max(ans, d_sum + len(t_set) * len(t_set))
    print(ans)


if __name__ == '__main__':
    solve()

Submission status: Accepted

Solution code:
import heapq
n, k = map(int, input().split())
sushi = []
for _ in range(n):
    t, d = map(int, input().split())
    sushi.append((d, t-1))

sushi.sort(reverse=True)
used = [False] * n
choose = []
rest = []
cost = 0
kind = 0
for i in range(n):
    d, t = sushi[i]
    if i < k:
        cost += d
        if used[t]:
            heapq.heappush(choose, d)
        else:
            used[t] = True
            kind += 1

    else:
        if not used[t]:
            heapq.heappush(rest, -d)
        used[t] = True


cost += kind ** 2
ans = cost

cnt = min(len(choose), len(rest))
for i in range(cnt):
    cost -= heapq.heappop(choose)
    cost += -heapq.heappop(rest)
    cost += 2 * kind + 1
    kind += 1
    ans = max(ans, cost)

print(ans)

Submission status: Accepted

Solution code:
import heapq
N, K = map(int,input().split())
L1 = []
L2 = []
S = set()
val = 0

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

L1.sort(reverse = True)

for i in range(K):
    val += L1[i][0]
    if L1[i][1] in S:
        heapq.heappush(L2, L1[i][0])
    else:
        S.add(L1[i][1])

ans = val + len(S) ** 2

for i in range(K, N):
    if len(L2) == 0:
        break
    if not L1[i][1] in S:
        S.add(L1[i][1])
        val -= heapq.heappop(L2)
        val += L1[i][0]
        ans = max(ans, val + len(S) ** 2)

print(ans)

Submission status: Accepted

Solution code:
def main():
    from collections import deque
    n, k, *td = map(int, open(0).read().split())
    s = list(zip(*[iter(td)] * 2))
    t = sorted(s, key=lambda x: x[1], reverse=True)

    variety = set()
    base = 0
    duplicated = deque([])

    for x, y in t[:k]:
        base += y
        if x in variety:
            duplicated.append(y)
        else:
            variety.add(x)

    now = base + len(variety) ** 2

    nx = deque(t[k:])
    c = len(variety)
    tmp = now
    while c < k and nx:
        i, j = nx.popleft()
        if i in variety:
            continue

        d = duplicated.pop()
        dif = j - d + 2 * c + 1
        variety.add(i)
        tmp = max(tmp, now + dif)
        now += dif
        c += 1

    print(tmp)


if __name__ == '__main__':
    main()

Submission status: Accepted

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:
        if ate[tds[i][0]]:
            i += 1
        else:
            ate[tds[i][0]] = True
            deliciousness -= d
            deliciousness += tds[i][1]
            variety += 1
            point = deliciousness + variety ** 2
            max_point = max(max_point, point)
            break

print(max_point)

Submission status: Accepted

Solution code:
def main():
    from heapq import heappush, heappop
    from operator import itemgetter
    import sys
    input = sys.stdin.readline

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

    dd = []
    for _ in range(N):
        t, d = map(int, input().split())
        t -= 1
        dd.append((d, t))
    dd.sort(key=itemgetter(0), reverse=True)

    ret = 0
    biggest = [-1] * N
    h = []
    kinds = 0
    for d, t in dd[:K]:
        ret += d
        if ~biggest[t]:
            if biggest[t] > d:
                heappush(h, d)
            else:
                heappush(h, biggest[t])
                biggest[t] = d
        else:
            kinds += 1
            ret += kinds * 2 - 1
            biggest[t] = d

    ans = ret
    for d, t in dd[K:]:
        if ~biggest[t]: continue
        if not h: break
        kinds += 1
        diff = d + kinds * 2 - 1 - h[0]
        ret += diff
        if ans < ret: ans = ret
        biggest[t] = d
        heappop(h)

    print(ans)


if __name__ == '__main__':
    main()

Submission status: Accepted

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())
        rem -= type_num
        for vals in sushis.values():
            for val in vals:
                rem_sushis.push(val)
        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())
        discard_num = type_num - k
        for _ in range(discard_num):
            eat_sushis.pop()
        for vals in sushis.values():
            for val in vals:
                rem_sushis.push(val)

    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 + eat_sushis.sum()

    print(ans)


if __name__ == '__main__':
    main()