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.

Solution (Python):
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)


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


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: 
            
            total += lst[j][0]
            kind += 1
            
            total -= lst2[n-1-i]
            ans = max(ans, total + kind * kind)
            flag = False
            j += 1
        else:
            j += 1
print (ans)

Solution (Python):
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])


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

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    
    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
        
        

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



print(max(ans_candidates.values()))

Solution (Python):
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() 

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

sushi_set = []
ans = 0
for p, t in sushis[:k]:
  
  ans += -p
  cnt_types[t] += 1
  sushi_set.append([p,t]) 
  
ans += len(cnt_types) ** 2 




max_types = []
max_types.append(ans)

for p, t in sushis[k:]: 
  if t in cnt_types:
    continue 
  else:
    
    

    while sushi_set:
      ps, ts = sushi_set.pop() 
      if cnt_types[ts] >= 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))

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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))

Solution (Python):

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
    
    kind[t]-=1
    cost-=d
    
    kind[nt]+=1
    cost+=nd
    K+=1
    ans=max(ans, cost+(K**2))

print(ans)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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])


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]
        
        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]
        
        remove_index -= 1
        if sushi_count[sushi[0]] > 1:
            kiso -= sushi[1]
            sushi_count[sushi[0]] -= 1
            break
    else:
        break
    
    ans = max(ans, kiso + now_kind * now_kind)
    now_kind += 1

print(ans)

Solution (Python):
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])


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])

    
ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1


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
        
        ans = max(ans, kiso + now_kind * now_kind)
        now_kind += 1
    else:
        continue

print(ans)

Solution (Python):
import sys

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)



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

print(ans)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
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)

Solution (Python):
import math
import copy
from collections import defaultdict
from collections import Counter
from collections import deque

from itertools import product

from itertools import permutations

from itertools import combinations
from bisect import bisect_left, bisect_right


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 = []


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


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)

Solution (Python):
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()

Solution (Python):
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))]))

Solution (Python):

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)

Solution (Python):









import sys
input = sys.stdin.readline
from heapq import heappush,heappop

N,K = map(int,input().split())
kind_value = [[] for _ in range(N+1)]
for _ in range(N):
  t,d = map(int,input().split())
  kind_value[t].append(d)
for li in kind_value:
  li.sort()
  
first_type = [(li[-1],i) for i,li in enumerate(kind_value) if li]
first_type.sort(reverse = True)

kind = 0
value_type1 = 0
second_type = []
second_type_sum = 0
second_type_cnt = 0

answer = 0
for v,i in first_type:
  kind += 1
  if kind > K:
    break
  value_type1 += v
  for d in kind_value[i][:-1]:
    heappush(second_type,d)
    second_type_sum += d
    second_type_cnt += 1
  while second_type_cnt > K-kind:
    d = heappop(second_type)
    second_type_cnt -= 1
    second_type_sum -= d
  value = kind*kind + value_type1 + second_type_sum
  if answer < value:
    answer = value

print(answer)

Solution (Python):

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


S = [[-float("inf")] for k in range(N)]
for k in range(N):
    t, d = map(int,input().split())
    S[t-1].append(d)





top = []
hoka = []
for k in range(N):
    S[k].sort()
    top.append(S[k].pop())
    hoka += S[k]

top.sort(reverse=True)
hoka.sort(reverse=True)



T = [0 for k in range(N+1)]
H = [0 for k in range(N+1)]
T[1], H[1] = top[0], hoka[0]
for k in range(2,N):
    T[k] = T[k-1] + top[k-1]
    H[k] = H[k-1] + hoka[k-1]

ans = 0
for k in range(1,K+1):
    ans = max(ans,T[k]+H[K-k]+k**2)
print(ans)

Solution (Python):
N, K = map(int,input().split())
S = [[-float("inf")] for k in range(N)]
for k in range(N):
    t, d = map(int,input().split())
    S[t-1].append(d)

top = []
hoka = []
for k in range(N):
    S[k].sort()
    top.append(S[k].pop())
    hoka += S[k]

top.sort(reverse=True)
hoka.sort(reverse=True)

T = [0 for k in range(N+1)]
H = [0 for k in range(N+1)]
T[1], H[1] = top[0], hoka[0]
for k in range(2,N):
    T[k] = T[k-1] + top[k-1]
    H[k] = H[k-1] + hoka[k-1]

ans = 0
for k in range(1,K+1):
    ans = max(ans,T[k]+H[K-k]+k**2)
print(ans)

Solution (Python):

import sys
from collections import defaultdict
from collections import deque


def solve(N: int, K: int, T: "List[int]", D: "List[int]"):
    ml = []

    for t,d in zip(T,D):
        ml.append((t,d))
    ml.sort(key=lambda x: x[1],reverse=True)

    mli = ml[0:K]
    mlo = deque(ml[K:])
    mli.reverse()

    tdic = defaultdict(int)
    for i in mli:
        tdic[i[0]] += 1
    
    max_score = sum(map(lambda x:x[1],mli)) + len(tdic)**2
    prev_score = max_score
    for i,p in enumerate(mli):
        if tdic[p[0]] > 1:
            while len(mlo) > 0:
                q = mlo.popleft()
                if q[0] not in tdic:
                    tdic[p[0]] -= 1
                    tdic[q[0]] += 1
                    mli[i] = q
                    tdic_len  = len(tdic)
                    new_score = prev_score - p[1] + q[1] - (tdic_len-1)**2 + (tdic_len)**2
                    max_score = max(max_score,new_score)
                    prev_score = new_score
                    break

    print(max_score)
    return



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

if __name__ == '__main__':
    main()

Solution (Python):

"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d
AC
"""
import sys

from collections import Counter

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

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

    rem = dict()                
    for t, d in sushi:
        if t not in pick_type:
            rem[t] = max(rem.get(t, 0), d)
    res = sorted([[k, v] for k, v in rem.items()], key=lambda x: x[1]) 

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


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


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

Solution (Python):

"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d

"""
import sys

from collections import Counter

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

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

    rem = dict()                
    for t, d in sushi:
        if t not in pick_type:
            rem[t] = max(rem.get(t, 0), d)
    res = sorted([[k, v] for k, v in rem.items()], key=lambda x: x[1]) 

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


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


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

Solution (Python):

import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


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

    q0 = []

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

    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

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

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

    return mmm


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

if __name__ == '__main__':
    main()

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


N, K = MI()

sushi = sorted([LI() for _ in range(N)], key=lambda x:x[1], reverse=True)

delicious_sushi = sushi[:K]

soudemonai_sushi = sushi[K:]

syurui_counter = Counter()
deliciousness = 0

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

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

f[len(syurui_counter)] = deliciousness






delicious_sushi = delicious_sushi[::-1]

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


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

print(max(f))

Solution (Python):
"""
漸化式的に解く
一般性を失わない範囲をちゃんと考える。
"置き換え可能"というのをどの部分で持たせられるかを考える
今回の場合は、リスト作りの際に作ったので十分
"""

from collections import defaultdict
N, K = map(int,input().split())

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

pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] 



cnt = 0
for t,d in pairs[:K]: 
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

n = len(kinds)
cnt += len(kinds)**2
ans = cnt

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    
    if len(change) == 0:
        break
    
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)

print(ans)

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

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

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

Solution (Python):
N,K=[int(a) for a in input().split(" ")]
TD = [[int(a) for a in input().split(" ")] for i in range(N)]
TD = sorted(TD, key=lambda arr: arr[1])[::-1]

types = set()
stack = []
ans, totald, i = 0, 0, 0
for i in range(K):
    t, d = TD[i]
    if t in types:
        stack.append(d)
    totald += d
    types.add(t)
ans = totald+len(types)**2

while len(stack)>0:
    while i<N and TD[i][0] in types:
        i+=1
    if i>=N: break
    t1, d1 = TD[i]
    lt = len(types)
    d2 = stack.pop()
    totald -= (d2-d1)
    types.add(t1)
    i+=1
    ans = max(ans, totald+len(types)**2)
print(ans)