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

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

    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)

            j -= 1
print(ans)

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

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

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

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

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





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


various = defaultdict(int)
L.sort(key = lambda i: i[1], reverse = True)

pool_der = []
ans = 0

for i in range(K):
    t, d = L[i]
    
    various[t] += 1
    
    if various[t] > 1:
        pool_der.append(d)
    ans += d

pool_der.sort(reverse = True)


now = len(various.items())
ans += now ** 2
L = L[K:]


pool_var = []
for i in L:
    
    
    if various[i[0]] == 0:
        pool_var.append(i[1])
        various[i[0]] += 1

pool_var.sort()


point = ans
M = min(len(pool_der), len(pool_var))

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

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

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

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

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

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

import heapq

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

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

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

ans[num_kinds] = tot + num_kinds**2

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

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

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

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

Solution (Python):
import sys
input = sys.stdin.readline
from collections import defaultdict

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

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

Solution (Python):



import sys

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

    for t, d in S:
        if t in st:
            
            y_0.append(y_0[-1]+d)
        else:
            st.add(t)
            y_1.append(y_1[-1]+d)
    

    max_sum = -1

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


if __name__ == '__main__':
    main()

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

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

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

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

Solution (Python):
from heapq import *

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

various = len(Top)


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

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

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

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


h_d = []

h_a = []

wk = 0
kind = 0

ap = set()
for i in range(K):
  d,t = I[i]
  wk += d
  if t not in ap:
    ap.add(t)
    kind += 1
  else:
    heapq.heappush(h_d, (d,t))

for i in range(K,N):
  d,t = I[i]
  if t not in ap:
    ap.add(t)
    
    heapq.heappush(h_a, ((-1) * d, t))
  
ans = wk + kind ** 2
while len(h_d) > 0 and len(h_a) > 0:
  d1,t1 = heapq.heappop(h_d)
  d2,t2 = heapq.heappop(h_a)
  d2 = (-1) * d2
  wk = wk - d1 + d2
  kind += 1
  ans = max(ans, wk + kind ** 2)
  
print(ans)

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

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


h_d = []

h_a = []

wk = 0
kind = 0

ap = set()
for i in range(K):
  d,t = I[i]
  wk += d
  if t not in ap:
    ap.add(t)
    kind += 1
  else:
    heapq.heappush(h_d, (d,t))

for i in range(K,N):
  d,t = I[i]
  if t not in ap:
    ap.add(t)
    
    heapq.heappush(h_a, ((-1) * d, t))
  
ans = wk + kind ** 2
while len(h_d) > 0 and len(h_a) > 0:
  d1,t1 = heapq.heappop(h_d)
  d2,t2 = heapq.heappop(h_a)
  d2 = (-1) * d2
  wk = wk - d1 + d2
  kind += 1
  ans = max(ans, wk + kind ** 2)
  
print(ans)

Solution (Python):
import heapq

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

eaten = [False]*N


d1, d2 = [0], [0]

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


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

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

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

Solution (Python):

import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

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

inf = float('inf')
mod = 1000000007




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

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

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

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

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

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

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

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

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



if __name__ == "__main__":
	main()

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

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

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


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

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

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

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

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

cur = 0
hq = []

for i in range(k):
    cur+=dset[i][0]
    for d in dset[i][1:]: heapq.heappush(hq, -d)


for dd in dset[k:]:
    for d in dd:
        heapq.heappush(hq, -d)

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

Solution (Python):

printn = lambda x: print(x,end='')
inn = lambda : int(input())
inl   = lambda: list(map(int, input().split()))
inm   = lambda:      map(int, input().split())
ins = lambda : input().strip()
DBG = True 
BIG = 10**18
R = 10**9 + 7

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

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

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

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


h = {}
for d,t in a:
    if t not in h:
        h[t] = d
    else:
        h[t] = max(h[t],d)
vals = []
for d,t in a:
    if t in h and d==h[t]:
        del h[t]
    else:
        vals.append(d)
vals.sort()


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

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

Solution (Python):
from collections import Counter
from heapq import heappush, heappop

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

Solution (Python):
import math

import queue
from collections import deque,defaultdict
import heapq as hpq
from sys import stdin,setrecursionlimit


ipt = stdin.readline
setrecursionlimit(10**7)

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

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

    print(ma)


    return

if __name__ == '__main__':
    main()

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

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

Solution (Python):

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

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

N_set = set()

sums = 0

cnt = [0 for i in range(N+1)]  
ng,ok,sub = [], [], []
heapq.heapify(ng)
heapq.heapify(ok)
heapq.heapify(sub)

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

ng = sub




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




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

print(ans)

Solution (Python):

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

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

N_set = set()

sums = 0

cnt = [0 for i in range(N+1)]  
ng,ok,sub = [], [], []
heapq.heapify(ng)
heapq.heapify(ok)
heapq.heapify(sub)

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

ng = sub




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




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

print(ans)

Solution (Python):


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

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

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

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

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

now = val + len(set_)**2


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

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

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

if '__main__' == __name__:
    resolve()

Solution (Python):
import sys
from collections import defaultdict
from heapq import heappush, heappop
from operator import itemgetter

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


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


if __name__ == '__main__':
    solve()

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

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

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


cost += kind ** 2
ans = cost

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

print(ans)

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

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

L1.sort(reverse = True)

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

ans = val + len(S) ** 2

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

print(ans)

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

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

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

    now = base + len(variety) ** 2

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

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

    print(tmp)


if __name__ == '__main__':
    main()

Solution (Python):
import sys
from operator import itemgetter
input = sys.stdin.buffer.readline

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

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

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

max_point = deliciousness + variety ** 2

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

print(max_point)

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

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

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

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

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

    print(ans)


if __name__ == '__main__':
    main()