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:
n, k = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(n)]

ab_s = sorted(ab, key=lambda x: x[1])
ab_r = list(reversed(ab_s))

t = [x[0] for x in ab_r]
d = [x[1] for x in ab_r]        

t_set = len(set(t[:k]))
d_sum = sum(d[:k])

res = [t_set**2 + d_sum]

from collections import Counter
t_cou = Counter(t[:k])

dp =[0]*(n+1)
j_old = k-1

for i in t_cou:
    dp[i] = t_cou[i]

for i in range(k,n):
    t_new = ab_r[i][0]
    if dp[t_new] ==0:
        for j in range(j_old,-1,-1):
            t_old = ab_r[j][0]
            if dp[t_old] <=1:
                pass
            else:
                dp[t_new] =1
                dp[t_old] -=1
                t_set +=1
                d_sum = d_sum - ab_r[j][1] + ab_r[i][1]
                res.append(t_set**2 + d_sum)
                j_old = j-1
                break
            j_old = 0

print(max(res))

Submission status: Accepted

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

ab_s = sorted(ab, key=lambda x: x[1])
ab_r = list(reversed(ab_s))

t = [x[0] for x in ab_r]
d = [x[1] for x in ab_r]        

t_set = len(set(t[:k]))
d_sum = sum(d[:k])

res = [t_set**2 + d_sum]

from collections import Counter
t_cou = Counter(t[:k])

dp =[0]*(n+1)
j_old = k-1

for i in t_cou:
    dp[i] = t_cou[i]

for i in range(k,n):
    t_new = ab_r[i][0]
    if dp[t_new] ==0:
        for j in range(j_old,-1,-1):
            t_old = ab_r[j][0]
            if dp[t_old] <=1:
                pass
            else:
                dp[t_new] =1
                dp[t_old] -=1
                t_set +=1
                d_sum = d_sum - ab_r[j][1] + ab_r[i][1]
                res.append(t_set**2 + d_sum)
                j_old = j-1
                break
            j_old = -1

print(max(res))

Submission status: Accepted

Solution code:
import sys
import queue
import collections

input_methods=['clipboard','file','key']
using_method=0
input_method=input_methods[using_method]

tin=lambda : map(int, input().split())
lin=lambda : list(tin())
mod=100000000777777

#+++++

def main():
	#a = int(input())
	n, k = tin()
	#s = input()
	al=[lin() for _ in range(n)]
	al.sort(key = lambda x : -x[1])
	use_t = set()
	t_max=collections.defaultdict(lambda :0)
	mm=[]
	ret_ug=0
	for t,v in al[:k]:
		ret_ug += v
		if t in use_t and t_max[t] >= v:
			mm.append(v)
		elif t in use_t:
			mm.append(t_max[t])
			t_max[t]=v
		else:
			t_max[t]=v
			use_t.add(t)
	
	nt=len(use_t)
	ret_ug += nt*nt
	
	nokori = collections.defaultdict(lambda :0)
	for t, v in al[k:]:
		if t in use_t:
			continue
		nokori[t]=max(nokori[t], v)
	upp=[]
	for k in nokori:
		upp.append(nokori[k])
	upp.sort(reverse = True)
	mm.sort()
	ret=ret_ug
	for up_v, mm_v in zip(upp, mm):
		ret_ug += up_v
		ret_ug -= mm_v
		ret_ug += 2*nt + 1
		nt += 1
		ret = max(ret, ret_ug)
		
	print(ret)
		
		
	
		
		
		
		
		
	
	
	
#+++++
isTest=False

def pa(v):
	if isTest:
		print(v)
		
def input_clipboard():
	import clipboard
	input_text=clipboard.get()
	input_l=input_text.splitlines()
	for l in input_l:
		yield l

if __name__ == "__main__":
	if sys.platform =='ios':
		if input_method==input_methods[0]:
			ic=input_clipboard()
			input = lambda : ic.__next__()
		elif input_method==input_methods[1]:
			sys.stdin=open('inputFile.txt')
		else:
			pass
		isTest=True
	else:
		pass
		#input = sys.stdin.readline
			
	ret = main()
	if ret is not None:
		print(ret)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
TD = [[int(j) for j in input().split()] for i in range(N)]
TD = sorted(TD,key=lambda x:-x[1])
point = 0
stack = []
kind = set()

for t,d in TD[:K]:
  point+=d
  if t in kind:
    stack.append(d)
  else:
    kind.add(t)

point+=len(kind)**2
ans = point

for t,d in TD[K:]:
  if not stack:
    break
  if not t in kind:
    point+=d-stack.pop()+2*len(kind)+1
    kind.add(t)
    ans=max(ans,point)

print(ans)

Submission status: Accepted

Solution code:
from math import ceil
from collections import defaultdict as dd
n, k = map(int, input().split())
sushi = []
for i in range(n):
    sushi.append(list(map(int, input().split())))
sushi.sort(key=lambda x: x[1])

dic = dd(list)
for sushi_type, oishisa in sushi:
    dic[sushi_type].append(oishisa)

sushi_type_num = len(dic)

#base point priority
dp = [0] * (sushi_type_num + 1)
eaten_sushi = dd(list)
base_point = 0
stack_remove = []
for i in range(k):
    sushi_type, oishisa = sushi.pop()
    eaten_sushi[sushi_type].append(oishisa)
    base_point += oishisa
    if len(eaten_sushi[sushi_type]) >= 2:
        stack_remove.append([sushi_type, oishisa])
bonus_point = len(eaten_sushi) ** 2
point = bonus_point + base_point
dp[len(eaten_sushi)] = point

#print(base_point, bonus_point, eaten_sushi, dp)

while sushi != []:
    #print(base_point, bonus_point, eaten_sushi, dp)
    sushi_type, oishisa = sushi.pop()
    if eaten_sushi[sushi_type] == []:
        #remove
        if stack_remove != []:
            remove_type, remove_oishisa = stack_remove.pop()
            eaten_sushi[remove_type].pop()
            base_point -= remove_oishisa
        #add
            eaten_sushi[sushi_type].append(oishisa)
            base_point += oishisa
            bonus_point = len(eaten_sushi) ** 2
            point = bonus_point + base_point
            dp[len(eaten_sushi)] = point
        else:
            break
print(max(dp))

Submission status: Accepted

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

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

# 各種類で最大のもの
P = []
# Pに入らなかったもの
Q = []

for k in sushi.keys():
    # 最大のものだけ
    sushi[k].sort(reverse=True)
    P.append(sushi[k][0])
    for x in sushi[k][1:]:
        Q.append(-x)

# heapqで最小値を扱う
heapq.heapify(P)
# heapqで最大値を扱う
heapq.heapify(Q)

cur = 0
# 全種類食べることが出来る場合
if K >= len(sushi.keys()):
    max_variation = len(sushi.keys())
    # 全種類について、最大のものを食べる
    cur += sum(P)
    # まだ食べられるので、Qを大きい方から食べる
    for _ in range(K - max_variation):
        cur += -heapq.heappop(Q)
    ans = cur + max_variation**2
# 全種類食べることが出来ない場合
else:
    # 最大値が大きい方からからK種類だけPに残す
    # 選べなかったものはQに突っ込む
    while len(P) > K:
        m = heapq.heappop(P)
        heapq.heappush(Q, -m)
    max_variation = K
    cur += sum(P)
    ans = cur + max_variation**2

# max_variation-1, ..., 1種類食べるとき
for i in range(max_variation - 1, 0, -1):
    # Pのうち、最小のものは一旦諦める
    p = heapq.heappop(P)
    cur -= p

    # 諦めたものをQに入れて、Qから最大のものを1つ食べる
    # 結果的にi種類ではなくなる場合もあるが、無視できる
    heapq.heappush(Q, -p)
    cur += -heapq.heappop(Q)
    ans = max(ans, cur + i**2)

print(ans)

Submission status: Accepted

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

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

# 各種類で最大のもの
P = []
# Pに入らなかったもの
Q = []

for k in sushi.keys():
    # 最大のものだけ
    sushi[k].sort(reverse=True)
    P.append(sushi[k][0])
    for x in sushi[k][1:]:
        Q.append(-x)

# heapqで最小値を扱う
heapq.heapify(P)
# heapqで最大値を扱う
heapq.heapify(Q)

cur = 0
# 全種類食べることが出来る場合
if K >= len(sushi.keys()):
    max_variation = len(sushi.keys())
    # 全種類について、最大のものを食べる
    cur += sum(P)
    # まだ食べられるので、Qを大きい方から食べる
    for _ in range(K - max_variation):
        cur += -heapq.heappop(Q)
    ans = cur + max_variation**2
# 全種類食べることが出来ない場合
else:
    # 最大値が大きい方からからK種類だけPに残す
    # 選べなかったものはQに突っ込む
    while len(P) > K:
        m = heapq.heappop(P)
        heapq.heappush(Q, -m)
    max_variation = K
    cur += sum(P)
    ans = cur + max_variation**2

# max_variation-1, ..., 1種類食べるとき
for i in range(max_variation - 1, 0, -1):
    # Pのうち、最小のものは一旦諦める
    p = heapq.heappop(P)
    cur -= p

    # 諦めたものをQに入れる
    heapq.heappush(Q, -p)
    q = -heapq.heappop(Q)
    cur += q
    ans = max(ans, cur + i**2)

print(ans)

Submission status: Accepted

Solution code:
import heapq
from sys import stdin

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

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = 0

if K <= len(sushi.keys()):
    max_variation = K
else:
    max_variation = len(sushi.keys())
    for _ in range(K-max_variation):
        cur -= heapq.heappop(Q)

cur += sum(P)
ans = cur + max_variation**2

for i in range(max_variation-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(max_variation-i-1)**2)

print(ans)

Submission status: Accepted

Solution code:
import heapq

n, k = map(int, input().split())
a = [0] * (n + 1)
typecnt = 0
h = []
heapq.heapify(h)
for _ in range(n):
    t, d = map(int, input().split())
    if a[t] == 0:
        typecnt += 1
        a[t] = 1
    heapq.heappush(h, [-d, t])
ans = 0
mintype = 0
m = []
heapq.heapify(m)
for _ in range(k):
    d, t = heapq.heappop(h)
    if a[t] == 1:
        mintype += 1
        a[t] = 0
    else:
        heapq.heappush(m, -d)
    ans -= d
cnt = ans
ans += mintype * mintype
for j in range(mintype, typecnt):
    if not m:
        break
    bonus = (j + 1) * (j + 1)
    while h:
        d, t = heapq.heappop(h)
        if not m:
            break
        if a[t] == 1:
            a[t] = 0
            l = heapq.heappop(m)
            cnt -= (d + l)
            ans = max(ans, cnt + bonus)
            break
print(ans)

Submission status: Accepted

Solution code:

from collections import defaultdict
from heapq import heappop, heappush

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

X.sort(key=lambda x: x[1])
pq = []
f_x = 0
ctr = defaultdict(int)
appeared = 0
for _ in range(K):
    t, d = X.pop()
    heappush(pq, (d, t))
    appeared += int(ctr[t] == 0)
    f_x += d
    ctr[t] += 1

res = f_x + appeared ** 2
while pq or X:
    while pq:
        d, t = heappop(pq)
        if ctr[t] > 1:
            ctr[t] -= 1
            f_x -= d
            break

    if not pq:
        break

    while X:
        t_x, d_x = X.pop()
        if ctr[t_x] == 0:
            heappush(pq, (d_x, t_x))
            f_x += d_x
            ctr[t_x] += 1
            appeared += 1
            res = max(res, f_x + appeared ** 2)
            break

print(res)

Submission status: Accepted

Solution code:
import heapq
from operator import itemgetter

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

sushi.sort(key=itemgetter(1), reverse=True)
que = []
kind = set()
res = 0
for t, d in sushi[:k]:
    if t in kind:
        heapq.heappush(que, d)
    else:
        kind.add(t)
    res += d
res += len(kind)**2

val = res
for t, d in sushi[k:]:
    if not que:
        break
    if t in kind:
        continue
    val += -heapq.heappop(que) + d + len(kind)*2 + 1
    kind.add(t)
    res = max(res, val)

print(res)

Submission status: Accepted

Solution code:
import sys
input=sys.stdin.readline
from heapq import heappop,heappush
hq=[]
n,k=map(int,input().split())
for i in range(n):
	t,d=map(int,input().split())
	heappush(hq,(-d,t))
D=0
dic={}
got=[]
for i in range(k):
	r=heappop(hq)
	got.append(r)
	if r[1] in dic:
		dic[r[1]]+=1
	else:
		dic[r[1]]=1
	D+=-r[0]
ans=D+len(dic)**2

while hq:
	while got:
		p=got.pop()
		if dic[p[1]]>1:
			dic[p[1]]-=1
			D-=-p[0]
			break
	if len(got)==0:
		print(ans)
		exit()

	r=heappop(hq)
	got.append(r)
	if r[1] in dic:
		dic[r[1]]+=1
	else:
		dic[r[1]]=1
	D+=-r[0]	
	ans=max(D+len(dic)**2,ans)
print(ans)

Submission status: Accepted

Solution code:
import sys
from heapq import *
input = sys.stdin.readline
N, K = map(int, input().split())
D, H, var, res, L = [], [], set(), 0, N.bit_length()+2
for _ in range(N):
    t,d = map(int, input().split())
    D.append((d<<L)+t)
D.sort(reverse=True)
for z in D[:K]:
    d, t = z>>L, z%(1<<L)
    if t not in var: var.add(t)
    else: heappush(H, d)
    res += d
x = len(var)
res += x**2
tmp = res
for z in D[K:]:
    d, t = z>>L, z%(1<<L)
    if H and t not in var:
        tmp += d-heappop(H)+2*x+1
        x += 1
        var.add(t)
        res = max(res,tmp)
print(res)

Submission status: Accepted

Solution code:
import sys
from heapq import *
input = sys.stdin.readline
N, K = map(int, input().split())
D, H, var, res = [], [], set(), 0
for _ in range(N):
    t,d = map(int, input().split())
    D.append((t,d))
D.sort(key=lambda x: x[1], reverse=True)
for t, d in D[:K]:
    if t not in var: var.add(t)
    else: heappush(H, d)
    res += d
x = len(var)
res += x**2
tmp = res
for t, d in D[K:]:
    if H and t not in var:
        tmp += d-heappop(H)+2*x+1
        x += 1
        var.add(t)
        res = max(res,tmp)
print(res)

Submission status: Accepted

Solution code:
# 種類が同じ寿司であれば、最も美味しさが高い寿司を選ぶべき。
# 従い、
# (1)各種類の最も美味しさが高い寿司グループ
# (2)それ以外の寿司グループ
# をいずれも降順に並べて、(1)から幾つ取るか(最低でも1つは取る)を全探索する
import sys
readline = sys.stdin.readline

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

from collections import defaultdict
sushi = defaultdict(list) # 種類・おいしさ　のdictionary

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

for value in sushi.values():
  value = sorted(value, reverse = True)
  best_sushi.append(value[0])
  other_sushi += value[1:]

best_sushi = sorted(best_sushi, reverse = True)
other_sushi = sorted(other_sushi, reverse = True)

best_num = 1 # best_sushiから選ぶ個数
other_num = K - 1 # other_sushiから選ぶ個数

if len(other_sushi) < other_num: # 選べない場合
  other_num = len(other_sushi)
  best_num = K - other_num
  
# それぞれの寿司から選ぶ個数の開始値が決定
# other_sushiの後ろの要素は切り離して良い

other_sushi = other_sushi[:other_num]
  
sushi_point = sum(best_sushi[:best_num]) + sum(other_sushi)
kind_point = best_num ** 2

ans = sushi_point + kind_point

# best寿司のインデックスを進めていく
for i in range(best_num, len(best_sushi)):
  sushi_point += best_sushi[i]
  if len(other_sushi) == 0:
    break
  rem = other_sushi.pop() # 取り除く寿司
  sushi_point -= rem
  kind_point = (i + 1) ** 2
  
  if ans < sushi_point + kind_point:
    ans = sushi_point + kind_point
    
print(ans)

Submission status: Accepted

Solution code:
# 種類が同じ寿司であれば、最も美味しさが高い寿司を選ぶべき。
# 従い、
# (1)各種類の最も美味しさが高い寿司グループ
# (2)それ以外の寿司グループ
# をいずれも降順に並べて、(1)から幾つ取るか(最低でも1つは取る)を全探索する
import sys
readline = sys.stdin.readline

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

from collections import defaultdict
sushi = defaultdict(list) # 種類・おいしさ　のdictionary

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

for value in sushi.values():
  value = sorted(value, reverse = True)
  best_sushi.append(value[0])
  other_sushi += value[1:]

best_sushi = sorted(best_sushi, reverse = True)
other_sushi = sorted(other_sushi, reverse = True)

#print("best_sushi",best_sushi)
#print("other_sushi",other_sushi)

best_num = 1 # best_sushiから選ぶ個数
other_num = K - 1 # other_sushiから選ぶ個数

if len(other_sushi) < other_num: # 選べない場合
  other_num = len(other_sushi)
  best_num = K - other_num
  
# それぞれの寿司から選ぶ個数の開始値が決定
# other_sushiの後ろの要素は切り離して良い

other_sushi = other_sushi[:other_num]
  
sushi_point = sum(best_sushi[:best_num]) + sum(other_sushi)
kind_point = best_num ** 2

ans = sushi_point + kind_point
#print("初期値",ans,"sum(best_sushi[:best_num])",sum(best_sushi[:best_num]),"sum(other_sushi)",sum(other_sushi),"best_num ** 2",best_num ** 2)

# best寿司のインデックスを進めていく
for i in range(best_num, len(best_sushi)):
  sushi_point += best_sushi[i]
  #print("----best_sushi[i]",best_sushi[i],"を追加")
  if len(other_sushi) == 0:
    break
  rem = other_sushi.pop() # 取り除く寿司
  sushi_point -= rem
  kind_point = (i + 1) ** 2
  
  #print("sushi_point",sushi_point,"kind_point",kind_point,"合計",sushi_point + kind_point)
  
  if ans < sushi_point + kind_point:
    ans = sushi_point + kind_point
    
print(ans)

Submission status: Accepted

Solution code:
import sys
def MI(): return map(int,sys.stdin.readline().rstrip().split())


N,K = MI()
dt = []
for i in range(N):
    t,d = MI()
    dt.append((d,t))

dt.sort(reverse=True)

basic_point = 0
ans = 0
used = set()
neta = [10**10]*(N+1)  # neta[t] = ネタtのある、最も左の位置
for i in range(K):
    d,t = dt[i]
    used.add(t)
    basic_point += d
    neta[t] = min(neta[t],i)
shurui = len(used)  # ネタの種類数
ans = basic_point + shurui**2

a,b = K-1,K
while a >= 0 and b < N:
    d1,t1 = dt[a]
    if neta[t1] != a:
        basic_point -= d1
        while b < N:
            d2,t2 = dt[b]
            if not t2 in used:
                used.add(t2)
                basic_point += d2
                shurui += 1
                ans = max(ans,basic_point + shurui**2)
                b += 1
                break
            else:
                b += 1
    a -= 1

print(ans)

Submission status: Accepted

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

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

D = defaultdict(list)
T = [0]*(N+1)
H = []
ans = [0]

for i, td in enumerate(TD):
  t, d = td[0], td[1]
  if i == K:
    break
  T[t] += 1
  D[t].append(d)
  ans[0] += d
  if T[t] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t] -= 1

cnt = sum(1 for c in T if c >= 1)
ans[0] += cnt**2

for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    if T[t1] == 0:
      d2, t2 = heappop(H)
      ans.append(ans[-1] + d1 - d2 + 2*cnt + 1)
      cnt += 1
      T[t1] += 1
  else:
    break

print(max(ans))

Submission status: Accepted

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

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

TD.sort(key = lambda x:x[1], reverse=True)
TD += [(TD[-1][0], 0)]

D = defaultdict(list)
T = [0]*N
H = []
ans = [0]

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans[0] += d

cnt = sum(1 for c in T if c >= 1)
ans[0] += cnt**2

for t, d in TD:
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0:
      ans.append(ans[-1] + d1 - d2 + 2*cnt + 1)
      cnt += 1
      T[t1-1] += 1
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

print(max(ans))
#print(ans)

Submission status: Accepted

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

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

TD.sort(key = lambda x:x[1], reverse=True)
TD += [(TD[-1][0], 0)]

D = defaultdict(list)
T = [0]*N
H = []
ans = 0

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans += d

cnt = sum(1 for c in T if c >= 1)
ans += cnt**2

for t in range(1, N+1):
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0 and d1 - d2 + 2*cnt + 1 >= 0:
      ans += d1 - d2 + 2*cnt + 1
      cnt += 1
      T[t1-1] += 1
      if T[t2-1] >= 2:
        T[t2-1] -= 1
      if T[t2-1] >= 2:
        heappush(H, (D[t2].pop(), t2))
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

#元のやつとの組み合わせ...？
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0
 
for t, d in TD:
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans, ans2))

Submission status: Accepted

Solution code:
import sys
from operator import itemgetter
from heapq import *
input = sys.stdin.readline



def main():
    n, k = map(int, input().split())
    sushi = [list(map(int, input().split())) for _ in range(n)]
    
    sushi.sort(key=itemgetter(1), reverse=True)
    
    key = []
    already = [False]*(n+1)
    variety = k
    taste = 0
    for i in range(k):
        v, t = sushi[i]
        taste += t
        if already[v]:
            variety -= 1
            heappush(key, t)
            continue
        already[v]= True
    
    ans = [variety**2+taste]
    for i in range(k, n):
        if variety == k:
            break
        v, t = sushi[i]
        if already[v]:
            continue
        already[v] = True
        taste += t-key[0]
        heappop(key)
        variety += 1
        ans.append(variety**2+taste)
        
    print(max(ans))
    
    
        
   
    
if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
def main():
    N, K = list(map(int, input().split()))
    sushi_list = list()
    for _ in range(N):
        t, d = list(map(int, input().split()))
        t -= 1
        sushi_list.append((t, d))
    sushi_list.sort(key=lambda s: s[1], reverse=True)  # おいしさの降順でソート
    # おいしさTop Kのすしをとるのを初期解とする
    neta_flag_list = [0] * N
    topk_extra_sushi_list = list()  # Top Kのすしで、各ネタ最大のおいしさはないすし
    additional_sushi_list = list()  # Top Kのすしのネタ以外で、各ネタの最大のおいしさのすし
    point_topk = 0
    for k, sushi in enumerate(sushi_list):
        if k < K:
            if neta_flag_list[sushi[0]] == 1:
                topk_extra_sushi_list.append(sushi)
            point_topk += sushi[1]
        else:
            if neta_flag_list[sushi[0]] == 0:
                additional_sushi_list.append(sushi)
        neta_flag_list[sushi[0]] = 1
    n_neta = K - len(topk_extra_sushi_list)
    point_topk += n_neta ** 2
    # その他のネタをt種類追加する場合を考えて、初期解と比較していく
    d_extra, d_additional = 0, 0
    ans = point_topk
    n_additional_neta = min(K - n_neta, len(topk_extra_sushi_list), len(additional_sushi_list))
    for t in range(1, n_additional_neta + 1):
        n_cur_neta = n_neta + t
        d_extra += topk_extra_sushi_list[-t][1]
        d_additional += additional_sushi_list[t - 1][1]
        point = point_topk - d_extra + d_additional + (n_cur_neta**2 - n_neta**2)
        ans = max(ans, point)
    print(ans)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            s.add(t)
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()

Submission status: Accepted

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

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

kind=set()

base=0
rest=[]

for i in range(k):
    t,d=td.pop()

    base+=d

    if t in kind:
        rest.append(d)
    
    kind.add(t)

cntk=len(kind)
ans=base+cntk*cntk

while len(rest)>0:
    bd=rest.pop()
    
    while len(td)>0:
        nt,nd=td.pop()
        if nt not in kind:
            break
    else:
        break

    base+=nd-bd
    cntk+=1
    ans=max(ans,base+cntk*cntk)
    kind.add(nt)

print(ans)

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
import math
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()
various = defaultdict(list)
que = [getList() for i in range(N)]

ans = 0
num = []
var_s = set()

# 美味しい順にK個とった時の幸福度
que.sort(reverse = True, key = lambda i: i[1])
for i in range(K):
    ans += que[i][1]
    # もし２番手以降ならあとで交換する用にとっておく
    if que[i][0] in var_s:
        num.append(que[i][1])
    var_s.add(que[i][0])

var = len(var_s)
ans += var ** 2

# 使ってない種類について各種類で一番大きさが大きいもの
left_l = defaultdict(int)
for i in range(N):
    if not que[i][0] in var_s:
        left_l[que[i][0]] = max(left_l[que[i][0]], que[i][1])

num.sort(reverse = True)
left_l = [i[1] for i in left_l.items()]
left_l.sort()

# M回交換する
opt = ans
M = min(len(num), len(left_l))
for i in range(M):
    u = num.pop()
    s = left_l.pop()
    # 寿司単体の幸福度
    opt -= (u - s)
    # 種類が増える分
    opt += 2 * var + 1
    var += 1
    ans = max(opt, ans)

print(ans)

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:
import sys
input = lambda : sys.stdin.readline().rstrip()
sys.setrecursionlimit(max(1000, 10**9))
write = lambda x: sys.stdout.write(x+"\n")


n,K = list(map(int, input().split()))
from collections import defaultdict
current = []
nums = defaultdict(int)
td = [None]*n
for i in range(n):
    t,d = list(map(int, input().split()))
    td[i] = (d,t)
td.sort(reverse=True)
ans = 0
for i in range(K):
    d,t = td[i]
    current.append((d, t))
    ans += d
    nums[t] += 1
j = K-1
total = len(nums)
ans += total**2
cur = ans
for i in range(K, n):
#     print(ans)
    d,t = td[i]
    if t not in nums:
        while j>=0 and nums[current[j][1]]<=1:
            j -= 1
        if nums[current[j][1]]>=2:
            nums[current[j][1]] -= 1
            nums[t] = 1
            cur -= current[j][0]
            cur += total*2+1 + d
            total += 1
            ans = max(ans, cur)
#             print(cur,j)
            j -= 1
print(ans)

Submission status: Accepted

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

def main():
  N, K=map(int,input().split())
  sushies=[list(map(int,input().split())) for i in range(N)]
  sushies.sort(reverse=True, key=lambda x:x[1])
  dsum=0
  tset=set()
  kubi=[]
  for i in range(K):
    dsum+=sushies[i][1]
    t=sushies[i][0]
    if t in tset:
      kubi.append(sushies[i])
    tset.add(t)
  x=len(tset)
  koho=[]
  for i in range(K,N):
    t=sushies[i][0]
    if not(t in tset):
      koho.append(sushies[i])
      tset.add(t)
  ind=0
  ans=dsum+x*x
  while(ind < len(kubi) and ind < len(koho)):
    dsum=dsum-kubi[-1-ind][1]+koho[ind][1]
    x+=1
    ans=max(ans, dsum+x*x)
    ind+=1
  print(ans)
  
if __name__=="__main__":
  main()