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())
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 d>top[t]:
            sub.append(top[t])
            top[t]=d
        else:
            sub.append(d)
top.sort(reverse=True)
top_sum=[0]
sub.sort(reverse=True)
sub_sum=[0]

for i in top:
    if i==0:
        top_sum.append(-1)
    else:
        top_sum.append(top_sum[-1]+i)
for i in sub:
    sub_sum.append(sub_sum[-1]+i)
ans=0
for i in range(1,k+1):
    if i+len(sub_sum)-1>=k and top_sum[i]!=-1:
        ans=max(ans,top_sum[i]+sub_sum[k-i]+i**2)
print(ans)

Solution (Python):
from collections import defaultdict

n, k = map(int, input().split())
top = [0] * n
sub = []

for i in range(n):
    t, d = map(int, input().split())
    t -= 1
    if top[t] == 0:
        top[t] = d
    else:
        if d > top[t]:
            sub.append(top[t])
            top[t] = d
        else:
            sub.append(d)

top.sort(reverse=True)
top_sum = [0]
sub.sort(reverse=True)
sub_sum = [0]

for p in top:
    if p == 0:
        top_sum.append(-1)
    else:
        top_sum.append(top_sum[-1] + p)

for p in sub:
    sub_sum.append(sub_sum[-1] + p)

ans = 0

for i in range(1, k + 1):
    if (len(sub_sum) - 1) >= (k - i) and top_sum[i] != -1:
        ans = max(ans, top_sum[i] + sub_sum[k - i] + i ** 2)

print(ans)

Solution (Python):
def main():
    N, K = map(int, input().split())
    SUSHI = []
    for _ in range(N):
        t, d = map(int, input().split())
        SUSHI.append((d, t))
    SUSHI.sort(reverse=True)

    NEW = {i: True for i in range(1, N+1)} 
    x, y = 0, 0 
    A, B = [], [] 
                  

    
    for d, t in SUSHI[:K]:
        if NEW[t]: 
            x += d
            y += 1 
            NEW[t] = False
        else:
            x += d
            A.append(d)
    for d, t in SUSHI[K:]:
        if NEW[t]:
            B.append(d)
            NEW[t] = False
    B.sort()
    
    
    

    answer = x + y**2
    while A and B: 
        a, b = A.pop(), B.pop() 
        x -= a
        x += b
        y += 1
        if x + y**2 > answer:
            answer = x + y**2
    print(answer)
main()

Solution (Python):
n,k=map(int,input().split())
INF=10**18
sushi=[[-INF] for _ in range(n)]
for _ in range(n):
  t,d=map(int,input().split())
  sushi[t-1].append(d)
for i in range(n):
  sushi[i].sort(reverse=True)
sushi.sort(key=lambda x:x[0],reverse=True)
q=[]
res=0
for i in range(k):
  res+=sushi[i][0]
  for j in sushi[i][1:]:
    q.append(j)
q.sort()
cur=res
res+=k**2
for var in range(1,k)[::-1]:
  a=q.pop()
  b=sushi[var][0]
  cur+=a-b
  res=max(res,cur+var**2)
print(res)

Solution (Python):
import heapq
n,k = map(int,input().split())
L = []
for i in range(n):
    t,d = map(int,input().split())
    L.append([d,t])
L.sort()
ans = []
cur = [0]*n
heap = []
var = 0
for i in range(k):
    if cur[L[n-i-1][1]-1] == 0:
        ans.append(L[n-i-1][0])
        cur[L[n-i-1][1]-1] = 1
        var += 1
    else:
        heapq.heappush(heap, L[n-i-1])
temp = 0
for i in range(len(ans)):
    temp += ans[i]
for i in range(len(heap)):
    temp += heap[i][0]
temp += var**2
now = temp
for i in range(n-k):
    if len(heap) != 0:
        a = heapq.heappop(heap)
        if cur[L[n-k-i-1][1]-1] == 0:
            cur[L[n-k-i-1][1]-1] = 1
            now -= var**2
            now -= a[0]
            var += 1
            now += var**2
            now += L[n-k-i-1][0]
            temp = max(temp, now)
        else:
            heapq.heappush(heap, a)
print(temp)

Solution (Python):
import heapq
n,k = map(int,input().split())
L = []
for i in range(n):
    t,d = map(int,input().split())
    L.append([d,t])
L.sort()
ans = []
cur = [0]*n
heap = []
var = 0
for i in range(k):
    if cur[L[n-i-1][1]-1] == 0:
        ans.append(L[n-i-1][0])
        cur[L[n-i-1][1]-1] = 1
        var += 1
    else:
        heapq.heappush(heap, L[n-i-1])
temp = 0
for i in range(len(ans)):
    temp += ans[i]
for i in range(len(heap)):
    temp += heap[i][0]
temp += var**2
now = temp
for i in range(n-k):
    if len(heap) != 0:
        a = heapq.heappop(heap)
        if cur[L[n-k-i-1][1]-1] == 0:
            ans.append(L[n-k-i-1][0])
            cur[L[n-k-i-1][1]-1] = 1
            now -= var**2
            now -= a[0]
            var += 1
            now += var**2
            now += L[n-k-i-1][0]
            temp = max(temp, now)
        else:
            heapq.heappush(heap, a)
print(temp)

Solution (Python):
from heapq import heappop,heappush
n,k=map(int,input().split())
s=[]
h=set()
l=[0]*n 
for _ in range(n):
    t,d=map(int,input().split())
    s.append((d,t-1))
s=sorted(s,reverse=True)
a=[]
num=0
for i in range(k):
    d,t=s[i]
    l[t]+=1
    h.add(t)
    num+=d
    if l[t]>1:
        heappush(a,(d,t))
ans=len(h)**2+num
for j in range(k,n):
    d,t=s[j]
    l[t]+=1
    if  l[t]>=2:
        continue 
    if len(a)==0:
        break
    d1,t1=heappop(a) 
    h.add(t)
    num=num-d1+d
    ans=max(ans,len(h)**2+num)
print(ans)

Solution (Python):

def main():
    N, K = map(int, input().split())
    l = []
    for _ in range(N):
        t, d = map(int, input().split())
        l.append((t, d))
    l.sort(key=lambda x:x[1], reverse=True)
    s = set()
    e = []
    t = 0
    for i in range(K):
        t += l[i][1]
        if l[i][0] in s:
            e.append(l[i])
        else:
            s.add(l[i][0])
    rl = [t + len(s) * len(s)]
    e.reverse()
    ei = 0
    for i in range(K, N):
        if l[i][0] in s:
            continue
        if ei >= len(e):
            break
        s.add(l[i][0])
        t -= e[ei][1]
        t += l[i][1]
        ei += 1
        rl.append(t + len(s) * len(s))
    return max(rl)
print(main())

Solution (Python):
def main():
    N, K = map(int, input().split())
    TD = [tuple(map(int, input().split())) for _ in range(N)]
    TD.sort(key=lambda x: x[1], reverse=True)
    s = set()
    rep = []
    for t, d in TD[:K]:
        if t in s:
            rep.append(d)
        else:
            s.add(t)
    m = sum(d for _, d in TD[:K])
    p = len(s) ** 2 + m
    for t, d in TD[K:]:
        if not rep:
            break
        if t in s:
            continue
        m += d - rep.pop()
        s.add(t)
        p = max(p, len(s) ** 2 + m)
    return p

print(main())

Solution (Python):
N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
M = len(set([t for t, d in TD]))
 
TD.sort(key=lambda x: x[1], reverse=True)
p, S, V = 0, [], set()
for t, d in TD[:K]:
    p += d
    if t in V:
        S.append(d)
    else:
        V.add(t)
p += len(V) ** 2
ans = p
for t, d in TD[K:]:
    if len(S) == 0:
        break
    if t not in V:
        p += 2 * len(V) + 1 - S.pop() + d
        V.add(t)
        ans = max(ans, p)
print(ans)

Solution (Python):
N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
M = len(set([t for t, d in TD]))

TD.sort(key=lambda x: x[1], reverse=True)
p, S, V = 0, [], set()
for t, d in TD[:K]:
    p += d
    if t in V:
        S.append(d)
    else:
        V.add(t)
p += len(V) ** 2
ans = p
for t, d in TD[K:]:
    if len(S) == 0:
        break
    if t not in V:
        p += 2 * len(V) + 1 - S.pop() + d
        V.add(t)
        ans = max(ans, p)
print(ans)

Solution (Python):
from collections import Counter
n,k = map(int,input().split())
td = [list(map(int,input().split())) for i in range(n)]
td.sort(key = lambda x:x[1],reverse=True)
ans = list(range(k))
ansv = sum((td[i][1] for i in range(k)))
c = Counter((td[i][0] for i in range(k)))
cnt = len(c)
ansv += len(c)**2
ptr = k
ansls = [ansv]
for i in range(k-1,-1,-1):
  t,d = td[ans[i]]
  if ptr < n and c[t] > 1:
    while ptr < n and c[td[ptr][0]] >= 1:
      ptr += 1
    if ptr >= n:
      break
    ansv = ansv-d+2*cnt+1+td[ptr][1]
    c[t] -= 1
    c[td[ptr][0]] += 1
    cnt += 1
    ptr += 1
    ansls.append(ansv)
print(max(ansls))

Solution (Python):
n,k=map(int,input().split())
X=[]
for i in range(n):
    t,d=map(int,input().split())
    X.append([t,d])
    
X.sort(key=lambda x:x[1],reverse=True)

V=set()
point=0
S=[]

for t,d in X[:k]:
    point+=d
    if t in V:
        S.append(d)
    else:
        V.add(t)
point+=len(V)**2

ans=point

for t,d in X[k:]:
    if not S:
        break
    if not t in V:
        point+=d-S.pop()+1+len(V)*2
        V.add(t)
        ans=max(ans,point)

print(ans)

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

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

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

used = [0 for _ in range(N)]
dup = []
deli = 0

for i in range(K):
    t, d = sushi[i]
    if used[t - 1]:
        dup.append(d)
    used[t - 1] = 1
    deli += d

vari = sum(used)
sat = vari**2 + deli

dup.sort(reverse=True)

for i in range(N - K):
    if not dup:
        break
    t, d = sushi[K + i]
    if used[t - 1]:
        continue
    used[t - 1] = 1
    deli -= dup.pop()
    deli += d
    vari += 1
    sat = max(sat, vari**2 + deli)

print(sat)

Solution (Python):


from heapq import heappop,heappush

INF = 10 ** 18

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

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


for i in range(N):
    S[i].sort(reverse = True)
    if not S[i]: S[i].append(-INF) 


S.sort(reverse = True)

uma = 0 
que = [] 
for i in range(K):
    
    uma += S[i][0]
    for j in range(1, len(S[i])):
        heappush(que, (-1) * S[i][j])


for i in range(K, N):
    for j in range(len(S[i])):
        heappush(que, (-1) * S[i][j])

ans = uma + K * K

for x in range(K - 1, 0, -1):
    v = S[x][0] 
    w = (-1) * que[0] 
    if v < w:
        heappop(que)
        heappush(que, (-1) * v)
        uma -= v; uma += w
    ans = max(ans, uma + x * x)

print(ans)

Solution (Python):


N, K = [int(x) for x in input().split()]
td = [0] * (N)
for i in range(N):
    td[i] = [int(x) for x in input().split()]



td = sorted(td, reverse = True)

var_p1 = [0] * N   
var_p0 = [0] * N   

n_p1 = 0
n_p0 = 0
for i in range(N):
    if i == 0 or (i > 0 and td[i][0] != td[i - 1][0]):  
        var_p1[n_p1] = td[i][1]  
        n_p1 += 1
    else:
        var_p0[n_p0] = td[i][1]  
        n_p0 += 1

var_p1.sort(reverse = True)
var_p0.sort(reverse = True)

val_taste = sum(var_p0[0 : K])
val_vars = 0
ans = 0

for j in range(1, K + 1):  
    val = var_p1[j - 1]
    if val <= 0: break  
    
    val_taste = val_taste + val - var_p0[K - j]
    val_vars = j * j
    ans = max(ans, val_taste + val_vars)

print(ans)

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


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

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

NEW = {i:True for i in range(1, n + 1)}
x, y = 0, 0
A, B = [], []

for d, t in SUSHI[:k]:
    if NEW[t]:
        x += d; y += 1
        NEW[t] = False
    else:
        x += d
        A.append(d)

for d, t in SUSHI[k:]:
    if NEW[t]:
        B.append(d)
        NEW[t] = False
B.sort()

answer = x + y**2
while A and B:
    a, b = A.pop(), B.pop()
    x -= a; x += b; y += 1
    if x + y**2 > answer:
        answer = x + y**2

print(answer)

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

N, K = map(int, input().split())
sushi = [None] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi[i] = (d, t)
sushi.sort(reverse=True)

types = set()
cand = []
s = x = 0

for d, t in sushi[:K]:
    if t in types:
        cand.append(d)
    else:
        types.add(t)
        x += 1
    s += d
cand.sort(reverse=True)

ans = s + x*x
for d, t in sushi[K:]:
    if t in types:
        continue
    if not cand:
        break
    dr = cand.pop()
    s += d - dr
    types.add(t)
    x += 1
    ans = max(ans, s + x*x)
print(ans)

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

N, K = map(int, input().split())
sushi = [None] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi[i] = (d, t-1)
sushi.sort(reverse=True)

types = set()
cand = []
s = x = 0

for d, t in sushi[:K]:
    if t in types:
        heappush(cand, d)
    else:
        types.add(t)
        x += 1
    s += d

ans = s + x*x
for d, t in sushi[K:]:
    if t in types:
        continue
    if not cand:
        break
    dr = heappop(cand)
    s += d - dr
    types.add(t)
    x += 1
    ans = max(ans, s + x*x)
print(ans)

Solution (Python):
import heapq
N,K = map(int,input().split())
td = [list(map(int,input().split())) for i in range(N)]
td.sort(key = lambda x:-x[1])
sumn = [[] for i in range(N)]
for i in range(N):
    heapq.heapify(sumn[i])
ans = 0
kind = {}
for t,d in td[:K]:
    if t not in kind:
        kind[t] = 1
    else:
        kind[t] += 1
    heapq.heappush(sumn[t-1],d)
kinds = len(kind)
ans = sum(sum(i) for i in sumn)
ans1 = ans +kinds**2
rem = td[K:]
nex = 0
candel = []
heapq.heapify(candel)
for k,v in kind.items():
    for i in range(v-1):
        heapq.heappush(candel,heapq.heappop(sumn[k-1]))
if candel:
    qd = heapq.heappop(candel)
    while K+nex < N:
        if td[K+nex][0] not in kind:
            kind[td[K+nex][0]] = 1
            kinds += 1
            ans = ans - qd + td[K+nex][1]
            ans1 = max(ans1,ans+kinds**2)
            nex += 1
            if candel:
                qd = heapq.heappop(candel)
            else:
                break
        else:
            nex += 1
print(ans1)

Solution (Python):
import sys

def solve():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    S = []
    for _ in range(N):
        t, d = map(int, input().split())
        S.append((d, t))
    S.sort(reverse = True)
    upperKCount = 0
    appeared = set()
    maxPoint = 0
    for i in range(K):
        maxPoint += S[i][0]
        if S[i][1] not in appeared:
            upperKCount += 1
            appeared |= {S[i][1]}
    maxPoint += upperKCount ** 2
    appeared = set()

    typeMax = []
    typeNum = 0
    typeLower = []
    others = 0
    for d, t in S:
        if t in appeared: 
            typeLower.append(d)
            others += 1
        else:
            typeMax.append(d)
            typeNum += 1
            appeared |= {t}
    typeMaxSum = [0] * (typeNum + 1)
    typeLowerSum = [0] * (others + 1)
    for i in range(typeNum): typeMaxSum[i+1] = typeMaxSum[i] + typeMax[i]
    for i in range(others): typeLowerSum[i+1] = typeLowerSum[i] + typeLower[i]
    for k in range(upperKCount + 1, min(typeNum, K) + 1):
        maxPoint = max(maxPoint, k ** 2 + typeMaxSum[k] + typeLowerSum[K - k])
    print(maxPoint)
    return 0

if __name__ == "__main__":
    solve()

Solution (Python):
N,K,*L=map(int,open(0).read().split());L=L[::-1]
I=set();X=0;U=[];D=[]
for d,t in sorted(zip(*[iter(L)]*2))[::-1]:
	if t in I:D+=d,
	elif X<K:I|={t};U+=d,;X+=1
D+=0,;i=K-X
while X and U[-1]+2*X-1<=D[i]:U.pop();X-=1;i+=1
print(sum(U+D[:i])+X**2)

Solution (Python):
N,K,*L=map(int,open(0).read().split())
S=sorted([(d,t)for t,d in zip(*[iter(L)]*2)],reverse=True)
In=set()
X=0
uni=[]
dub=[]
for d,t in S:
	if t not in In:
		In.add(t)
		uni.append(d)
		X+=1
		if X==K:
			break
	else:
		dub.append(d)
det=sum(dub[:K-X])
dub=dub[K-X:]
for d in dub:
	if uni[-1]+2*X-1<=d:
		uni.pop()
		X-=1
		det+=d
	else:
		break
print(sum(uni)+det+X**2)

Solution (Python):
N,K,*L=map(int,open(0).read().split())
S=sorted([(d,t)for t,d in zip(*[iter(L)]*2)],reverse=True)
In=set()
C=X=0
uni=[]
dub=[]
for d,t in S:
	if t not in In:
		In.add(t)
		uni.append(d)
		X+=1
		if X==K:
			break
	else:
		dub.append(d)
det=0
C=X
for d in dub:
	if C<K:
		det+=d
		C+=1
	elif uni[-1]+2*X-1<=d:
		uni.pop()
		X-=1
		det+=d
print(sum(uni)+det+X**2)

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], reverse=True)

types = set()
can_remove = []
d_sm = 0

for t, d in td[:k]:
	d_sm += d
	if t in types:
		can_remove.append(d)
	else:
		types.add(t)

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

left = td[k:][::-1]

while can_remove and left:
	t, d = left.pop()
	if t not in types:
		types.add(t)
		d_rmv = can_remove.pop()
		d_sm += d - d_rmv
		type_cnt += 1
		ans = max(ans, d_sm + type_cnt ** 2)

print(ans)

Solution (Python):
from heapq import heappush, heappop

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

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

types = set()
can_remove = []
d_sm = 0

for t, d in td[:k]:
	d_sm += d
	if t in types:
		can_remove.append(d)
	else:
		types.add(t)

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

left = []
for t, d in td[k:]:
	d = -d
	heappush(left, [d, t])

while can_remove and left:
	d, t = heappop(left)
	d = -d
	if t not in types:
		types.add(t)
		d_rmv = can_remove.pop()
		d_sm += d - d_rmv
		type_cnt += 1
		ans = max(ans, d_sm + type_cnt ** 2)

print(ans)

Solution (Python):
N,K=map(int,input().split())
sushi=[]
for i in range(N):
    t,d=map(int,input().split())
    sushi.append((t,d))
sushi.sort(key=lambda x:(x[0],-x[1]))

X=[sushi[0][1]]
Y=[]
for i in range(1,N):
    if sushi[i-1][0]==sushi[i][0]:
        Y.append(sushi[i][1])
    else:
        X.append(sushi[i][1])
X.sort(reverse=True)
Y.sort(reverse=True)
SX=[0]
SY=[0]
for i in X:
    SX.append(SX[-1]+i)
for i in Y:
    SY.append(SY[-1]+i)
ans=0




for x in range(max(1,K-len(Y)),min(len(SX),K+1)):
    ans=max(SX[x]+SY[K-x]+x*x,ans)
print(ans)

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

def main():
    N,K = map(int,input().split())
    sushi = []
    for _ in range(N):
        t,d = map(int,input().split())
        sushi.append((-d,t))
    sushi.sort(key=itemgetter(0))
    use,only,unuse = [],[],sushi[K:]
    neta = defaultdict(int)
    x = 0
    for i in range(K):
        d,t = sushi[i]
        if neta[t] == 0:
            only.append((-d,t))
        else:
            use.append((-d,t))
        neta[t] += 1
        x -= d
    
    l = len(neta)
    heapq.heapify(use)
    heapq.heapify(only)
    heapq.heapify(unuse)
    
    ans = x + l**2
    while (len(use) != 0 and len(unuse) != 0 and l < K):
        cd,ct = heapq.heappop(unuse)
        cd *= -1
        if neta[ct] != 0:
            continue
        heapq.heappush(only,(cd,ct))
        rd,rt = heapq.heappop(use)
        x += cd-rd
        l += 1
        neta[ct] += 1
        neta[rt] -= 1
        ans = max(ans,x+l**2)
        
    print(ans)

if __name__ == "__main__":
    main()

Solution (Python):
from heapq import heappush, heappop
N, K, *L = map(int, open(0).read().split())
ls = []
for t,d in zip(*[iter(L)]*2):
  ls += [(d,t)]
used = []
cnt = [0]*(N+1)
hq = []
S = set()
ls.sort(reverse=True)
m = 0
for i in range(K):
  d,t = ls[i]
  cnt[t] += 1
  m += d
  if cnt[t]>1:
    used.append(d)
  S.add(t)
used.sort(reverse=True)
m += len(S)**2
for i in range(K,N):
  d,t = ls[i]
  if t not in S:
    heappush(hq,(-d,t))
ans = [m]
while used:
  d = used.pop()
  while hq:
    if hq[0][1] not in S:
      break
    heappop(hq)
  else:
    break
  m -= d
  n,s = heappop(hq)
  m += -n+2*len(S)+1
  ans.append(m)
  S.add(s)
print(max(ans))

Solution (Python):
from heapq import heappush, heappop
N, K, *L = map(int, open(0).read().split())
ls = []
for t,d in zip(*[iter(L)]*2):
  ls += [(d,t)]
used = []
cnt = {}
hq = []
S = set()
ls.sort(reverse=True)
m = 0
for i in range(K):
  d,t = ls[i]
  cnt[t] = cnt.get(t,0)+1
  m += d
  if cnt[t]>1:
    heappush(used,(d,t))
  S.add(t)
m += len(S)**2
for i in range(K,N):
  d,t = ls[i]
  if t not in S:
    heappush(hq,(-d,t))
ans = m
while used:
  d,t = heappop(used)
  while hq:
    if hq[0][1] not in S:
      break
    heappop(hq)
  else:
    break
  m -= d
  n,s = heappop(hq)
  m += -n+2*len(S)+1
  ans = max(ans,m)
  S.add(s)
print(ans)

Solution (Python):
import sys
import heapq as h

input=sys.stdin.readline

N,K=map(int,input().split())
sushi=[[0,0] for i in range(0,N)]
for i in range(0,N):
    t,k=map(int,input().split())
    sushi[i][1]=t
    sushi[i][0]=k

sushikind=[[] for i in range(0,N+1)]
for i in range(0,N):
    sushikind[sushi[i][1]].append(sushi[i][0])

kindmax=[0 for i in range(0,N+1)]

for i in range(0,N):
    if sushikind[i]:
        kindmax[i]=max(sushikind[i])

sushi.sort(reverse=True)
sum=0
kind=[]
kind=set(kind)
extra=[]
flag=[0 for i in range(0,N+1)]
for i in range(0,K):
    sum+=sushi[i][0]
    kind.add(sushi[i][1])
    if sushi[i][0]==kindmax[sushi[i][1]]:
        if flag[sushi[i][1]]==1:
            extra.append(sushi[i][0])
        else:
            flag[sushi[i][1]]=1
    else:
        extra.append(sushi[i][0])


set={i for i in range(1,N+1)}
set-=kind
appendix=[kindmax[i] for i in set if kindmax[i]!=0 ]
extra.sort()
appendix.sort(reverse=True)
k=len(kind)
first=k
ans=k**2+sum
for i in range(k,min(K,k+len(appendix))):
    k+=1
    sum+=appendix[k-first-1]-extra[k-first-1]
    test=k**2+sum
    if test>ans:
        ans=test

print(ans)

Solution (Python):

n,k= map(int, input().split())
t= [list(map(int, input().split())) for i in range(n)]
q=[]
ans=0
r=[0]*(n+1)

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

for i in range(k):
    ans += t[i][1]
    if r[t[i][0]]==0:
        r[t[i][0]]=1
    else:
        q.append(t[i][1])
v=sum(r)
ans+=v**2
cnt=ans


for i in range(k,n):
    if r[t[i][0]]==0 and len(q)!=0:
        r[t[i][0]]=1
        x=q.pop()
        cnt+=t[i][1]+2*v+1-x
        v+=1
        ans=max(cnt,ans)
print(ans)

Solution (Python):
import collections
from collections import deque
n,k= map(int, input().split())
t= [list(map(int, input().split())) for i in range(n)]
q=deque()
ans=0
r=[0]*(n+1)

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

for i in range(k):
    ans += t[i][1]
    if r[t[i][0]]==0:
        r[t[i][0]]=1
    else:
        q.append(t[i][1])
v=sum(r)
ans+=v**2
cnt=ans


for i in range(k,n):
    if r[t[i][0]]==0 and len(q)!=0:
        r[t[i][0]]=1
        x=q.pop()
        cnt+=t[i][1]+2*v+1-x
        v+=1
        ans=max(cnt,ans)
print(ans)

Solution (Python):
from collections import defaultdict as dd
from heapq import heappop, heappush

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q1 = dd(list)
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

ans = score + kind*kind
while q2 and q3:
    score -= heappop(q2)+heappop(q3)
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Solution (Python):
from collections import defaultdict as dd
from heapq import heappop, heappush

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q1 = dd(list)
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

ans = score + kind*kind
while q2 and q3:
    d1 = heappop(q2)
    d2 = heappop(q3)
    score -= d1+d2
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Solution (Python):
from collections import defaultdict as dd
from heapq import heappop, heappush

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

ans = score + kind*kind

q1 = dd(list)
for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

while q2 and q3:
    d1 = heappop(q2)
    d1 = -d1

    d2 = heappop(q3)
    score -= d2
    score += d1
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Solution (Python):
n,m = map(int,input().split())
a = []
for i in range(n):
    r,t = map(int,input().split())
    a.append((t,r))
a.sort(reverse = True)

d = {}
cd = []
dp = 0
x = 0
 
for i in range(m):
    z,c = a[i]
    dp += z
    if c in d:
        cd.append(z)
    else:
        d[c] = z
        x += 1

cd.sort(reverse = True)
ans = dp + x**2
c = ans
 
for i in range(m,n):
    if len(cd) == 0:
        break
    z,v = a[i]
    if v in d:
        continue
    d[v] = z
    x += 1
    c = c - cd.pop()+z+2*x-1
    if c > ans:
        ans = c
print(ans)

Solution (Python):
import sys
def input():
    return sys.stdin.readline()[:-1]
N,K=map(int, input().split())
s=[None]*N
for i in range(N):
    t,d=map(int,input().split())
    s[i]=(d,t)
s.sort(reverse=True)
l=[0]*N
d=0
for i in s[:K]:
    l[i[1]-1]+=1
    d+=i[0]
c=0
for i in l:
    if i:
        c+=1
a=d+c**2
h=K-1
for i in s[K:]:
    if not l[i[1]-1]:
        while h>=0:
            if l[s[h][1]-1]>1:
                l[i[1]-1]=1
                l[s[h][1]-1]-=1
                c+=1
                d=d+i[0]-s[h][0]
                a=max(a,d+c**2)
                h-=1
                break
            h-=1
        else:
            break
print(a)

Solution (Python):
import sys

sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))

N, K = lr()
TD = [lr() for _ in range(N)]
TD.sort(key=lambda x: x[1], reverse=True)
used = set()
overlap = []
for t, d in TD[:K]:
    if t in used:
        overlap.append(d)
    else:
        used.add(t)

delicious = sum([x[1] for x in TD[:K]])
answer = delicious + len(used) ** 2
for i in range(K, N):
    if not overlap:
        break
    if TD[i][0] in used:
        continue
    delicious += (TD[i][1] -  overlap.pop())
    used.add(TD[i][0])
    if answer < delicious + len(used)**2:
        answer = delicious + len(used)**2

print(answer)

Solution (Python):
import heapq

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

ans = 0
t_set = set()
candidates = []
for d, t in dt[:K]:
  if t in t_set:
    heapq.heappush(candidates, [d, t])
    ans += d
  else:
    t_set.add(t)
    ans += d

x = len(t_set)
ans += x**2

if len(candidates) == 0:
  print(ans)
  quit()

ans = [ans]
for d, t in dt[K:]:
  if t in t_set:
    continue
  cand_d, cand_t = heapq.heappop(candidates)
  _ans = ans[-1] - cand_d + d + (2*x+1)
  x += 1
  ans.append(_ans)
  t_set.add(t)
  if len(candidates) == 0:
    break

print(max(ans))

Solution (Python):
n,k=map(int,input().split())
l=[]
for i in range(n):
    l.append(list(map(int,input().split())))
l.sort()
l.append([-1,-1])
a=[]
b=[]
for i in range(n):
    if l[i][0]==l[i+1][0]:
        b.append(l[i][1])
    else:
        a.append(l[i][1])
ans=len(a)*len(a)+sum(a)

if len(b)>0:
    ans+=sum(b)
    b.sort()
    a.sort()
    i=0
    ai=0
    bi=0
    an=len(a)
    bn=len(b)
    while i<n-k:
        if ai<an and bi<bn:
            ap=a[ai]
            bp=b[bi]
            if ap+2*(an-ai-1)+1>=bp:
                bi+=1
                ans-=bp
            else:
                ans-=ap+2*(an-ai-1)+1
                ai+=1
        elif ai<an:
            ap=a[ai]
            ans-=ap+2*(an-ai-1)+1
            ai+=1
        elif bi<bn:
            ans-=b[bi]
            bi+=1
        i+=1
    print(ans)
else:
    a.sort(reverse=True)
    print(sum(a[:k])+k*k)

Solution (Python):
n,k=map(int,input().split())
l=[]
for i in range(n):
    l.append(list(map(int,input().split())))
l.sort(key=lambda x: x[0]*1000000000000+x[1])
l.append([-1,-1])
a=[]
b=[]
for i in range(n):
    if l[i][0]==l[i+1][0]:
        b.append(l[i][1])
    else:
        a.append(l[i][1])
ans=len(a)*len(a)+sum(a)

if len(b)>0:
    ans+=sum(b)
    b.sort()
    a.sort()
    i=0
    ai=0
    bi=0
    an=len(a)
    bn=len(b)
    while i<n-k:
        if ai<an and bi<bn:
            ap=a[ai]
            bp=b[bi]
            if ap+2*(an-ai-1)+1>=bp:
                bi+=1
                ans-=bp
            else:
                ans-=ap+2*(an-ai-1)+1
                ai+=1
        elif ai<an:
            ap=a[ai]
            ans-=ap+2*(an-ai-1)+1
            ai+=1
        elif bi<bn:
            ans-=b[bi]
            bi+=1
        i+=1
    print(ans)
else:
    a.sort(reverse=True)
    print(sum(a[:k])+k*k)

Solution (Python):
n,k=map(int,input().split())
inf=10**18
sushi=[[-inf] for _ in range(n)]
for i in range(n):
  t,d=map(int,input().split())
  sushi[t-1].append(d)
for i in range(n):
  sushi[i].sort(reverse=True)
sushi.sort(key=lambda x:x[0],reverse=True)
q=[]
res=0
for i in range(k):
  res+=sushi[i][0]
  for j in sushi[i][1:]:
    q.append(j)
q.sort()
cur=res
res+=k**2
for var in range(1,k)[::-1]:
  a=q.pop()
  b=sushi[var][0]
  cur+=a-b
  res=max(res,cur+var**2)
print(res)

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

list0 = []
list1 = []
set1 = set()

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

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

for t, d in td:
    if t in set1:
        list0.append(d)
    else:
        set1.add(t)
        list1.append(d)

list0.sort(reverse=True)
list1.sort(reverse=True)

various = min(K, len(list1))
selected0 = []
selected1 = []
for i in range(various):
    selected1.append(list1.pop(0))
if K > various:
    for i in range(K - various):
        selected0.append(list0.pop(0))

sum0 = sum(selected0)
sum1 = sum(selected1)
ans = sum0 + sum1 + various ** 2
while list0 and list0[0] > selected1[-1]:
    sum0 += list0.pop(0)
    sum1 -= selected1.pop(-1)
    various -= 1
    tmp = sum0 + sum1 + various ** 2
    if tmp <= ans:
        break
    ans = tmp

print(ans)