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:
from collections import defaultdict
import sys

input = sys.stdin.readline

n, k = map(int, input().split())
e = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    e[t].append(d)

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)

Submission status: Accepted

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

input = sys.stdin.readline

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

e.sort(key=itemgetter(0), reverse=True)

b = e[:k]  # 暫定選択（おいしさ基礎ポイント降順取得）

used = [False] * n
hq = []

ans = 0  # 暫定選択（おいしさ基礎ポイント降順取得）
tn = 0  # 種類数
for d, t in b:
    if not used[t]:
        used[t] = True
        tn += 1
    else:
        heappush(hq, (d, t))
    ans += d
ans += pow(tn, 2)  # 種類ボーナス

c = e[k:]  # 取得しなかった予備

# print(b)
# print(c)
# print(hq)
# print(used)
s = ans
cur = 0
while hq:
    d, t = heappop(hq)
    s -= d
    changed = False
    while cur < n - k:
        cd, ct = c[cur]
        cur += 1
        if not used[ct]:
            used[ct] = True
            changed = True
            s += cd
            s += tn * 2 + 1
            tn += 1
            if s > ans:
                ans = s
            break
    if not changed:
        break
print(ans)

Submission status: Accepted

Solution code:
import heapq

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

total = 0
types = set()
remain = []
for i in range(K):
    d, t = heapq.heappop(sushi)
    total -= d
    if t not in types:
        types.add(t)
    else:
        # 余りの寿司はremainにappend
        heapq.heappush(remain, -d)

types_len = len(types)
ans = total + types_len ** 2

while sushi and remain:
    d, t = heapq.heappop(sushi)
    if t not in types:
        types.add(t)
        remain_d = heapq.heappop(remain)
        total = total - d - remain_d
        types_len += 1
        ans = max(ans, total + types_len ** 2)

print(ans)

Submission status: Accepted

Solution code:
from operator import itemgetter
n,k = map(int,input().split())
li=[]
s=[]
t=[]
g=0
ans=[]
retu=[0]*n
p=0
q=0
for _ in range(n):
    a,b=map(int,input().split())
    li.append([a,b])
li=sorted(li,key=itemgetter(1),reverse=True)
for i in range(k):
    g+=li[i][1]
    if retu[li[i][0]-1]==0:
        retu[li[i][0]-1]+=1
        p+=1
    else:
        t.append(li[i][1])
        q+=1
g+=p**2
ans.append(g)
for i in range(k,n):
    if q==0:
        break
    if retu[li[i][0]-1]==0:
        g=g-t[q-1]+li[i][1]+2*p+1
        p+=1
        q-=1
        ans.append(g)
        retu[li[i][0]-1]+=1
print(max(ans))

Submission status: Accepted

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

sushi = []

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

sushi.sort(key=lambda A : A[1], reverse=True)

# 美味しいものから順にK個選ぶ
nowSushi = []
sumTaste = 0
numKind = 0
take = [0] * N  # 種類は高々N

for i in range(K) :
    t, d = sushi[i]
    nowSushi.append((t,d))
    sumTaste += d
    if take[t-1] == 0 :
        numKind += 1
    take[t-1] += 1

addSushi = []  # 種類を増やせる寿司
for i in range(K,N) :
    t, d = sushi[i]
    if take[t-1] > 0 :  # すでに選択済み
        continue
    addSushi.append((t,d))
    take[t-1] += 1

ans = sumTaste + numKind**2
L = len(addSushi)
changeCount = 0

while nowSushi :
    removeSushiT, removeSushiD = nowSushi.pop()
    if take[removeSushiT-1] <= 1 :
        continue
    if changeCount >= L :
        break

    addSushiT, addSushiD = addSushi[changeCount]
    changeCount += 1

    sumTaste = sumTaste - removeSushiD + addSushiD
    numKind = numKind + 1

    if (sumTaste + numKind**2) > ans :
        ans = sumTaste + numKind**2

print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)
#'''


def caluculate_score():
	return num_kind * num_kind + taste
	
	
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
table = [[] for _ in range(N)]
eat_kind = [0] * N
num_kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	table[i] = (d[i], t[i] - 1)
	
table.sort()

eat = []
taste = 0

for _ in range(K):
	tmp = table.pop()
	taste += tmp[0]
	eat_kind[tmp[1]] += 1
	if eat_kind[tmp[1]] == 1:
		num_kind += 1
	eat.append(tmp)
	
ans = caluculate_score()
ind_eat = K - 1
ind_table = N - K - 1

while ind_eat >= 0 and ind_table >= 0:
	while ind_eat >= 0 and eat_kind[eat[ind_eat][1]] <= 1:
		ind_eat -= 1
		
	while ind_table >= 0 and eat_kind[table[ind_table][1]] >= 1:
		ind_table -= 1
		
	if not(ind_eat >= 0 and ind_table >= 0):
		break

	eat_kind[eat[ind_eat][1]] -= 1
	eat_kind[table[ind_table][1]] += 1
	num_kind += 1
	taste += table[ind_table][0] - eat[ind_eat][0]
	eat[ind_eat], table[ind_table] = table[ind_table], eat[ind_eat]
	score = caluculate_score()
	if score > ans:
		ans = score
		
	ind_eat -= 1
	ind_table -= 1

print(ans)

Submission status: Accepted

Solution code:
# Greedyを一ひねり
# 最初はGreedyに取る。それから種類を増やす。2個目以降を削除して、新たなネタの1個目を加える。
# あるネタの中でポイントが最高かどうかが重要。これを　True　/　False　で管理する
n, k = list(map(int, input().split()))

sushi = [list(map(int, input().split())) + [None] for _ in range(n)]
sushi.sort(key = lambda s: s[1] , reverse = True) #値段の降順

# あるネタの中でポイントが最高なら、 Trueそうでなければ、　False
neta_seen = set()
for i in range(n):
	if sushi[i][0] in neta_seen:
		sushi[i][2] = False
	else:
		sushi[i][2] = True
		neta_seen.add(sushi[i][0])

# n_kind = len([s[1] for s in sushi[:k] if s[1] == True ]) 
# なんかうまくいかないので、最初からk番目までのTrueの数を愚直に書こう

n_kind = 0
for i in range(k):
	if sushi[i][2]:
		n_kind += 1

ans = sum([s[1] for s in sushi[:k]]) + n_kind ** 2
#print(sushi, ans)

now = ans
adopt = k
trash = k-1

while(True):

	while(adopt < n and sushi[adopt][2] == False):
		adopt += 1
	while(trash >= 0 and sushi[trash][2] == True):
		trash -= 1
	if not(adopt < n and trash >= 0):
		break

	now = now + sushi[adopt][1] - sushi[trash][1]
	now = now - n_kind**2 + (n_kind+1) ** 2

#	print(adopt, trash, now)

	n_kind += 1
	adopt += 1
	trash -= 1

	if now > ans:
		ans = now

print(ans)

Submission status: Accepted

Solution code:
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
total = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    total -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, -d)

length = len(group)
ans = length**2 + total

while td and td_:
    d, t = heapq.heappop(td)

    if t not in group:
        group.add(t)
        d_ = heapq.heappop(td_)
        total = total - d_ - d
        length += 1
        ans = max(ans, length**2 + total)
print(ans)

Submission status: Accepted

Solution code:
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
total = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    total -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, -d)

ans = len(group)**2 + total

while td and td_:
    d, t = heapq.heappop(td)

    if t not in group:
        group.add(t)
        d_ = heapq.heappop(td_)
        total = total - d_ - d
        ans = max(ans, len(group)**2 + total)
print(ans)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD += [[0, 0]]
TD.sort(key = lambda x: x[1], reverse = True)
TD.sort(key = lambda x: x[0]) #ネタの種類、おいしさを、ネタの種類でソート

lst = []
for i in range(1, N + 1):
    if TD[i-1][0] != TD[i][0]:
        lst += [[TD[i][1], 1]]
    else:
        lst += [[TD[i][1], 0]]

lst.sort(key = lambda x: x[1], reverse = True)
lst.sort(key = lambda x: x[0], reverse = True)
# print (lst)

total = 0
kind = 0
lst2 = []
for i in range(K):
    total += lst[i][0]
    if lst[i][1] == 1:
        kind += 1
    else:
        lst2 += [lst[i][0]]

ans = total + kind * kind
# print (total, ans)
# print (lst2)
n = len(lst2)
j = K
for i in range(n):
    flag = True
    while flag:
        if j >= N:
            print (ans)
            exit()
        if lst[j][1] == 1: #ネタの種類が増えるものだけ考える
            # print (lst[j])
            total += lst[j][0]
            kind += 1
            # print (lst2[n-1-i])
            total -= lst2[n-1-i]
            ans = max(ans, total + kind * kind)
            flag = False
            j += 1
        else:
            j += 1
print (ans)

Submission status: Accepted

Solution code:
from collections import defaultdict

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

sushi.sort(key=lambda x: -x[1])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    # print("t, d = ", t, d)
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop()
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x
        # print("remove", remove_candidates)
        # print(x, umami, ans_candidates[x])

    if i == k - 1:
        x = len(kinds.keys())
        ans_candidates[x] = umami + x * x
        remove_candidates.sort(key=lambda x: -x)
        # print(umami)
        # print(kinds, x)
        # print(remove_candidates)
        # print(x, umami, ans_candidates[x])

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))

Submission status: Accepted

Solution code:
from collections import Counter

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

sushis = [[0 for _ in range(2)] for _ in range(n)]
cnt_sushi_type = {}
for i in range(n):
  t,p = map(int, input().split())
  sushis[i][0] = -p #ポイント、大きいものから取り出したいのでマイナスにして入れる
  sushis[i][1] = t #種類

cnt_types = Counter() #寿司の種類覚えておく用のdict

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

sushi_set = []
ans = 0
for p, t in sushis[:k]:
  #大きいものからk個選んで得点を計算
  ans += -p
  cnt_types[t] += 1
  sushi_set.append([p,t]) #集合を保存
  
ans += len(cnt_types) ** 2 #ネタ数＊ネタ数を加点

##ここまで貪欲選択


max_types = []
max_types.append(ans)
#max_types[len(cnt_types)] = ans
for p, t in sushis[k:]: #これまで選んでないものを選択した場合についてみていく
  if t in cnt_types:
    continue #選んでない種類のネタを選びたいので、すでに選んでるネタであれば何もせず飛ばす
  else:
    #選んだ寿司集合にある一番得点の低い寿司を、現在見ている新ネタの寿司と入れ替える
    #実際に入れ替えるとめんどくさい（計算量的にappendとかしんどい）ので、ansとcnt_typesをいじることで調整する

    while sushi_set:
      ps, ts = sushi_set.pop() #小さい順で取り出す
      if cnt_types[ts] >= 2:
        #そのネタが2個以上あるなら一個減らして新セットを作る
        tmp =  max_types[-1]- (-ps) + (-p) + 2*len(cnt_types) + 1
        max_types.append(tmp)
        cnt_types[t] += 1
        cnt_types[ts] -= 1
        break #次のまだ見ぬネタを探しに行く
        
#各種類選んだ最大値から最大のものを選択で終了
print(max(max_types))

Submission status: Accepted

Solution code:
from heapq import *

N, M = list(map(int, input().split()))
sushi = []
for _ in range(N):
    t, d = list(map(int, input().split()))
    sushi.append([t, d])
sushi.sort(key=lambda x: x[1], reverse=True)

num_sushi = {}
tmp = 0
f = [0] * (N+1)
h = []
for i in range(M):
    s = sushi[i]
    if s[0] in num_sushi:
        num_sushi[s[0]] += 1
    else:
        num_sushi[s[0]] = 1
    heappush(h, (s[1], s[0]))
    tmp += s[1]

f[len(num_sushi)] = tmp

for i in range(M, N):
    if sushi[i][0] in num_sushi:
        continue
    while len(h) > 0:
        s = heappop(h)
        if num_sushi[s[1]] >= 2:
            num_sushi[s[1]] -= 1
            heappush(h, (sushi[i][1], sushi[i][0]))
            num_sushi[sushi[i][0]] = 1
            f[len(num_sushi)] = f[len(num_sushi)-1] + sushi[i][1] - s[0]
            break

ans = 0
for i, g in enumerate(f):
    if g <= 0: continue
    if g + i*i > ans:
        ans = g + i*i
print(ans)

Submission status: Accepted

Solution code:
import sys
import math 
ni = lambda: int(ns())
na = lambda: list(map(int, input().split()))
ns = lambda: input()
N,K = na()
S = [na() for i in range(N)]
S.sort(key=lambda x:x[1],reverse=True)
res = []
tmp=0
d=set([])
ds = []
for t,dd in S[:K]:
    tmp+=dd
    if t in d:
        ds.append(dd)
    else:
        d.add(t)
ans = tmp + len(d)**2
ax = ans

for t,dd in S[K:]:
    if len(ds)==0:
        break
    if t not in d:
        ax +=2*len(d)+1
        ax += dd-ds.pop()
        d.add(t)
        ans = max(ax,ans)
print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict

def getlist():
	return list(map(int, input().split()))
D = defaultdict(int)
V = []

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

V = sorted(V)

ans = []
ex = []
k = 0
for i in range(N):
	x = V[-1 - i][1]
	if k < K:
		if D[x] == 0:
			ans.append(V[-1 - i])
			D[x] += 1
			k += 1
		else:
			ex.append(V[-1 - i])
	else:
		ex.append(V[-1 - i])


if k < K:
	for i in range(K - k):
		ans.append(ex[i])
	eex = []
	for i in range(K - k, len(ex)):
		eex.append(ex[i])
else:
	eex = ex
ans = sorted(ans)
eex = sorted(eex)


A = 0
for i in range(K):
	A += ans[i][0]
A += k * k
n = min(k, N - K)
for i in range(n):
	if ans[i][0] + ((k - i) ** 2) - ((k - i - 1) ** 2) < eex[-1 - i][0]:
		A += eex[-1 - i][0] - ans[i][0] - ((k - i) ** 2) + ((k - i - 1) ** 2)
	else:
		break

print(A)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
td = [list(map(int,input().split())) for _ in range(N)]
td.sort(key=lambda x:-x[1])
l,r = td[:K],td[K:]
used_l = [0]*N
for i in range(K):
    used_l[l[i][0]-1] += 1

remove = []
for i in range(K-1,-1,-1):
    if used_l[l[i][0]-1] >= 2:
        remove.append(l[i][1])
        used_l[l[i][0]-1] -= 1

used_r = used_l.copy()
add = []
for i in range(N-K):
    if used_r[r[i][0]-1] == 0:
        add.append(r[i][1])
        used_r[r[i][0]-1] += 1

t = sum(used_l)

ans = t**2
for i in range(K):
    ans += l[i][1]

ref = ans
for i in range(min(len(remove),len(add))):
    ref += (t+1)**2-t**2+add[i]-remove[i]
    t += 1
    ans = max(ans,ref)

print(ans)

Submission status: Accepted

Solution code:
import heapq
if __name__ == "__main__":
    n,k = map(int,input().split())
    lst = [list(map(int,input().split())) for i in range(n)]
    selected = [False] * (n+1)
    pq = []
    lst.sort(key=lambda x:x[1], reverse=True)
    f = [0] * (n+1)
    cnt = 0
    res = 0
    for i in range(k):
        res += lst[i][1]
        if selected[lst[i][0]]:
            heapq.heappush(pq, lst[i][1])
        else:
            selected[lst[i][0]]=True
            cnt += 1
    f[cnt] = res
    now = k
    while len(pq)!=0 and now<n:
        diff=0
        while now < n and selected[lst[now][0]]:
            now +=1
        if now >= n:
            break
        diff -= heapq.heappop(pq)
        diff += lst[now][1]
        selected[lst[now][0]] = True
        cnt += 1
        f[cnt]=f[cnt-1]+diff
    res = 0 
    for i in range(n+1):
        if f[i]!=0 and f[i]+i**2>res:
            res=f[i]+i**2
    print("{}".format(res))

Submission status: Accepted

Solution code:
#coding:utf-8
from Queue import PriorityQueue
from collections import defaultdict
pq=PriorityQueue()
N,K=map(int,raw_input().split())
T=[];D=[];I=range(N)
for i in range(N):
    t,d=map(int,raw_input().split())
    T.append(t);D.append(d)
    pq.put((-d,t,i))    
P=zip(D,T,I) 
P.sort();P.reverse()
cost=0
kind=defaultdict(int)
curPq=PriorityQueue()
for d,t,i in P[:K]:
    kind[t]+=1
    cost+=d
    curPq.put((d,t,i))
ans=cost+(len(kind)**2)
K=len(kind)
while True:
    d,t,i=None,None,None
    while not curPq.empty():
        d,t,i=curPq.get()
        if kind[t]>1:break
    else: break
    nd,nt,ni=None,None,None
    while not pq.empty():
        nd,nt,ni=pq.get()
        nd=-nd
        if kind[nt]==0:break
    else: break
    # delete: (d,t,i)
    kind[t]-=1
    cost-=d
    # add: (nd,nt,ni)
    kind[nt]+=1
    cost+=nd
    K+=1
    ans=max(ans, cost+(K**2))

print(ans)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
sushi=[]
for i in range(n):
    sushi.append([])
for i in range(n):
    t,d=map(int,input().split())
    sushi[t-1].append(d)
Max=[]
Rest=[]
v=0
for i in sushi:
    if not i==[]:
        Max.append(max(i))
        i.remove(max(i))
        for j in i:
            Rest.append(j)
        v=v+1
Max.sort(reverse=True)
Rest.sort(reverse=True)
Max_s=[]
Rest_s=[0]
Sum=0
for i in Max:
    Sum=Sum+i
    Max_s.append(Sum)
Sum=0
l_r=0
for i in Rest:
    Sum=Sum+i
    Rest_s.append(Sum)
    l_r=l_r+1
ans=0
for i in range(1,min(k,v)+1):
    if k-i<=l_r:
        temp=Max_s[i-1]+Rest_s[k-i]+i**2
        if ans<temp:
            ans=temp
print(ans)

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
td=[list(map(int,input().split())) for _ in range(N)]
td=[[td[i][0]-1,td[i][1]] for i in range(N)]

hozon_a=[[] for _ in range(N)]
hozon_b=[[] for _ in range(N)]

td.sort(key=lambda x:x[1],reverse=True)
for i in range(K):
    hozon_a[td[i][0]].append(td[i][1])
for i in range(K,N):
    if hozon_a[td[i][0]]==[] and hozon_b[td[i][0]]==[]:
        hozon_b[td[i][0]].append(td[i][1])

cnt=len([None for a in hozon_a if a!=[]])
ans=sum([sum(hozon_a[i]) for i in range(N)])+cnt**2

hozon_aa=[hozon_a[i][1:] for i in range(N) if len(hozon_a[i])>1]
hozon_aa=[e for inner_list in hozon_aa for e in inner_list]
hozon_bb=[hozon_b[i][0] for i in range(N) if hozon_b[i]!=[]]

hozon_aa.sort(reverse=True)
hozon_bb.sort(reverse=False)

cnt1=cnt
temp=ans
for i in range(min(K-cnt,len(hozon_aa),len(hozon_bb))):
    temp+=hozon_bb.pop()-hozon_aa.pop()+2*cnt1+1
    cnt1+=1
    if ans<temp:
        ans=temp
print(ans)

Submission status: Accepted

Solution code:
from collections import deque
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)

s = set()
bests = []
rems = deque()
for d, t in dt:
    if t in s:
        rems.append(d)
    else:
        s.add(t)
        bests.append(d)
    if len(bests) == K:
        break
x = len(bests)
k = x
p = sum(bests) + x * x
while k < K:
    p += rems.popleft()
    k += 1
maxp = p
while len(rems) and rems[0] > bests[-1]:
    p -= bests.pop() + x * x
    x -= 1
    p += rems.popleft() + x * x
    maxp = max(maxp, p)
print(maxp)

Submission status: Accepted

Solution code:
n, k = 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)
sushi = [[] for i in range(2)]
used = set()
for t, d in td:
    if t in used:
        sushi[0].append(d)
    else:
        sushi[1].append(d)
        used.add(t)
y = k - min(k, len(sushi[0]))
sushi = list(reversed(sushi[0][:k])) + sushi[1][:k]
sum = sum(sushi[:k])
ans = sum + y ** 2
for i in range(len(sushi) - k):
    sum += sushi[k + i] - sushi[i]
    ans = max(ans, sum + (y + i + 1) ** 2)
print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
N, K = map(int, input().split())

sushi_array = []

sushi_count = [-1] * N
sushi_kind = 0
for i in range(N):
    t, d = map(int, input().split())
    sushi_array.append([t - 1, d])
    if sushi_count[t - 1] == -1:
        sushi_count[t - 1] = 0
        sushi_kind += 1

sushi_array = sorted(sushi_array, key=lambda x: -x[1])
# print(sushi_array, sushi_kind)

now_kind = 0
kiso = 0
for sushi in sushi_array[:K]:
    kiso += sushi[1]
    if sushi_count[sushi[0]] == 0:
        now_kind += 1
    sushi_count[sushi[0]] += 1

ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1

now_kind += 1
while now_kind <= sushi_kind:

    while add_index < N:
        sushi = sushi_array[add_index]
        # print(sushi)
        add_index += 1
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break

    else:
        break
    while remove_index >= 0:
        sushi = sushi_array[remove_index]
        # print(sushi)
        remove_index -= 1
        if sushi_count[sushi[0]] > 1:
            kiso -= sushi[1]
            sushi_count[sushi[0]] -= 1
            break
    else:
        break
    # print(ans,kiso,now_kind * now_kind)
    ans = max(ans, kiso + now_kind * now_kind)
    now_kind += 1

print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
N, K = map(int, input().split())

sushi_array = []

sushi_count = [-1] * N
sushi_kind = 0
for i in range(N):
    t, d = map(int, input().split())
    sushi_array.append([t - 1, d])
    if sushi_count[t - 1] == -1:
        sushi_count[t - 1] = 0
        sushi_kind += 1

sushi_array = sorted(sushi_array, key=lambda x: -x[1])
# print(sushi_array, sushi_kind)

now_kind = 0
kiso = 0
can_remove = []
for sushi in sushi_array[:K]:
    kiso += sushi[1]
    if sushi_count[sushi[0]] == 0:
        now_kind += 1
        sushi_count[sushi[0]] = 1
    else:
        can_remove.append(sushi[1])

    # print(sushi_count)
ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1

# print(ans, kiso, now_kind)
now_kind += 1

for sushi in sushi_array[add_index:]:
    if len(can_remove) == 0 or now_kind > sushi_kind:
        break
    if sushi_count[sushi[0]] == 0:
        kiso += sushi[1] - can_remove.pop()
        sushi_count[sushi[0]] = 1
        # print(ans, kiso, now_kind)
        ans = max(ans, kiso + now_kind * now_kind)
        now_kind += 1
    else:
        continue

print(ans)

Submission status: Accepted

Solution code:
import sys
# input = sys.stdin.readline
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]

table = [[] for _ in range(N)]
for i, (ti, di) in enumerate(td):
    table[ti - 1].append((i, ti, di))

tops = []
others = []
for row in table:
    if len(row) > 0:
        for i, item in enumerate(sorted(row, key=lambda x: -x[2])):
            if i == 0: tops.append(item)
            else: others.append(item)

tops.sort(key=lambda x: -x[2])
others.sort(key=lambda x: -x[2])
x_max = len(tops)

tops_sumtab = [0]
tmp = 0
for i, ti, di in tops:
    tmp += di
    tops_sumtab.append(tmp)
    
others_sumtab = [0]
tmp = 0
for i, ti, di in others:
    tmp += di
    others_sumtab.append(tmp)

# print("#", tops)
# print("#", others)
ans = 0
for i in range(K + 1):
    if K - i > len(others): continue
    elif i > len(tops): break
    # print("##", tops_sumtab[i], others_sumtab[K - i], i ** 2)
    ans_tmp = tops_sumtab[i] + others_sumtab[K - i] + i ** 2
    if ans < ans_tmp: ans = ans_tmp

print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        x.append(di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
idx = 1
for i in range(n - k - 1, -1, -1):
    if idx == len(x) + 1:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - x[-idx]
        p += 1
        idx += 1
        ans = max(ans, z + p ** 2)

print(ans)

Submission status: Accepted

Solution code:
import heapq

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        x.append(di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
j = 1
for i in range(n - k - 1, -1, -1):
    if j == len(x) + 1:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - x[-j]
        p += 1
        j += 1
        ans = max(ans, z + p ** 2)

print(ans)

Submission status: Accepted

Solution code:
import heapq

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        heapq.heappush(x, di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
for i in range(n - k - 1, -1, -1):
    if len(x) == 0:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - heapq.heappop(x)
        p += 1
        ans = max(ans, z + p ** 2)

print(ans)

Submission status: Accepted

Solution code:
from heapq import heappop
from heapq import heappush
from heapq import heappushpop

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

a=[list(map(int,input().split()))[::-1] for i in range(N) ]
a.sort(reverse=True)

l=[]
dic1={}
dic2={}

for i in range(K):
    [d,t]=a[i]
    if t in dic1:
        if d>dic2[t]:
            heappush(l,dic2[t])
            dic2[t]=d
        else: heappush(l,d)

    else: 
        dic1[t]=1
        dic2[t]=d


s=sum(l)+sum(dic2.values())
x=len(dic1)
ans=s+x**2

for i in range(K,N):
    if l==[]: break
        
    [d,t]=a[i]
    if not t in dic1:
        dic1[t]=1
        s+=d-heappop(l)
        x+=1
        ans=max(ans,s+x**2) 

    

print(ans)

Submission status: Accepted

Solution code:
import math
import copy
from collections import defaultdict
from collections import Counter
from collections import deque
# 直積 A={a, b, c}, B={d, e}:のとき，A×B={(a,d),(a,e),(b,d),(b,e),(c,d),(c,e)}: product(A, B)
from itertools import product
# 階乗 P!: permutations(seq), 順列 {}_len(seq) P_n: permutations(seq, n)
from itertools import permutations
# 組み合わせ {}_len(seq) C_n: combinations(seq, n)
from itertools import combinations
from bisect import bisect_left, bisect_right
# import numpy as np

def inside(y, x, H, W):
    return 0 <= y < H and 0 <= x < W

# 四方向: 右, 下, 左, 上
dy = [0, -1, 0, 1]
dx = [1, 0, -1, 0]

def i_inpl(): return int(input())
def s_inpl(): return map(int,input().split())
def l_inpl(): return list(map(int, input().split()))
INF = float("inf")
MAX_DIGIT = 50

############
############
############

N, K = l_inpl()
td = []
for _ in range(N):
    ti, di = l_inpl()
    td.append([ti, di])

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

x = 0 # 美味しさの和
kind = set()
changable = []

# 美味しい順からソートして上位K個を採用する
for i in range(K):
    x += td[i][1]
    if td[i][0] not in kind: # 新しいネタの種類のとき
        kind.add(td[i][0])
    else: # 過去に追加したネタの種類のとき
        changable.append(td[i][1])
    
ans = x + len(kind) ** 2

# 採用した上位K個を元に下位を採用したときのスコアを計算する
for i in range(K, N):
    # ありえるネタの種類と各種類において最大のスコアを採用し終わった
    if len(changable) == 0:
        break
    # 新しいネタがすでに存在している種類の場合
    if td[i][0] in kind:
        continue
    kind.add(td[i][0])
    x = x + td[i][1] - changable.pop()
    tmp = x + len(kind) ** 2
    ans = max(ans, tmp)

print(ans)

Submission status: Accepted

Solution code:
import sys
def LI(): return [int(x) for x in sys.stdin.readline().split()]

def main():
    N, K = LI()
    TD = []
    for _ in range(N):
        TD.append(LI())

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

    from collections import Counter
    import heapq

    counter = Counter()
    S = TD[:K]
    TD = TD[K:]
    queue = []
    for t, d in S:
        counter[t] += 1
        heapq.heappush(queue, (d, t))

    sum_ = sum([i[1] for i in S])
    maximum = sum_ + len(counter)**2
    while len(queue) > 0:
        d, t = heapq.heappop(queue)
        if counter[t] == 1:
            continue
        counter[t] -= 1
        sum_ -= d
        while len(TD) > 0:
            new_t, new_d = TD.pop(0)
            if counter[new_t] == 0:
                counter[new_t] += 1
                sum_ += new_d
                maximum = max(maximum, sum_ + len(counter)**2)
                break

    print(maximum)


main()

Submission status: Accepted

Solution code:
import sys
from itertools import accumulate
N, K = map(int, input().split())
dt = [tuple(map(int, sys.stdin.readline().split()))[::-1] for _ in range(N)]
dt.sort()
ans = sum([d for d, t in dt[N-K:]])
repet = []
seen = set()
for _ in range(K):
    d, t = dt.pop()
    if t not in seen:
        seen.add(t)
    else:
        repet.append(d)
sp = len(seen)
repet.append(0)
uniq = [0]
for _ in range(N-K):
    d, t = dt.pop()
    if t not in seen:
        uniq.append(d)
        seen.add(t)
repet = list(accumulate(repet[::-1]))
uniq = list(accumulate(uniq))
print(ans + max([(sp + i)**2 - r + u for i, (r, u) in enumerate(zip(repet, uniq))]))

Submission status: Accepted

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


top = [0] * N
sub = []
for _ in range(N):
    t, d = map(int, input().split())
    t -= 1
    if top[t] == 0:
        top[t] = d
    else:
        if top[t] >= d:
            sub.append(d)
        else:
            sub.append(top[t])
            top[t] = d
top.sort(reverse=True)
sub.sort(reverse=True)


res = [0]
for i in range(N):
    if top[i] != 0:
        res.append(res[-1] + top[i])
    else:
        res.append(-1)
w = [0]
for i in range(len(sub)):
    w.append(w[-1] + sub[i])


ans = 0
for i in range(1, K+1):
    if (len(w) - 1 >= K-i) and (res[i] != -1):
        ans = max(res[i] + w[K-i] + i**2, ans)
print(ans)