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):
import sys
input = sys.stdin.readline

N, K = map(int, input().split())
nums = [[] for i in range(N)]
for i in range(N):
    t, d = map(int, input().split())
    t -= 1
    nums[t].append(d)
maxs = []
rest = []
for i in range(N):
    nums[i].sort(reverse=True)
    if len(nums[i]) > 0:
        maxs.append(nums[i][0])
        for n in nums[i][1:]:
            rest.append(n)
        
maxs.sort(reverse=True)
if len(maxs) > K:
    maxs = maxs[:K]
    for n in maxs[K:]:
        rest.append(n)
    
rest.sort(reverse=True)
other = []
idx = 0
while len(maxs)+len(other) < K:
    other.append(rest[idx])
    idx += 1
rest = rest[idx:]

types = len(maxs)


for i,r in enumerate(rest):
    diff = types*types - pow(types-1, 2)
    if diff < r-maxs[len(maxs)-1-i]:
        maxs[len(maxs)-1-i] = r
        types -= 1
    else:
        break
ans = types*types
for num in maxs:
    ans += num
for num in other:
    ans += num
print(ans)

Solution (Python):
N, K = map(int, input().split())
nums = [[] for i in range(N)]
for i in range(N):
    t, d = map(int, input().split())
    t -= 1
    nums[t].append(d)
maxs = []
rest = []
for i in range(N):
    nums[i].sort(reverse=True)
    if len(nums[i]) > 0:
        maxs.append(nums[i][0])
        for n in nums[i][1:]:
            rest.append(n)
        
maxs.sort(reverse=True)
if len(maxs) > K:
    maxs = maxs[:K]
    for n in maxs[K:]:
        rest.append(n)
    
rest.sort(reverse=True)
other = []
idx = 0
while len(maxs)+len(other) < K:
    other.append(rest[idx])
    idx += 1
rest = rest[idx:]

types = len(maxs)


for i,r in enumerate(rest):
    diff = types*types - pow(types-1, 2)
    if diff < r-maxs[len(maxs)-1-i]:
        maxs[len(maxs)-1-i] = r
        types -= 1
    else:
        break
ans = types*types
for num in maxs:
    ans += num
for num in other:
    ans += num
print(ans)

Solution (Python):

import sys
try:
    from typing import List
except ImportError:
    pass


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    ss = sorted(zip(d, t), reverse=True)
    gs = []
    mst = set()
    for di, ti in ss[:K]:
        if ti in mst:
            gs.append((di, ti))
        else:
            mst.add(ti)
    dm = {}
    for di, ti in ss[K:]:
        if ti in dm or ti in mst:
            continue
        dm[ti] = di
    hs = sorted(((di, ti) for ti, di in dm.items()), reverse=True)
    cd = sum(di for di, ti in ss[:K])
    ct = len(mst)
    m = cd + ct * ct
    for (gdi, gti), (hdi, hti) in zip(reversed(gs), hs):
        cd += hdi - gdi
        ct += 1
        m = max(m, cd + ct * ct)
    print(m)


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):
N,K,*L = map(int, open(0).read().split())
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  ls.append((d,t))
ls.sort()
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = ls.pop()
  ans += d
  if t not in S:
    num += 1
    S.add(t)
  else:
    pre.append(d)
pre.sort(reverse=True)
ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = ls.pop()
  if t in S:
    continue
  S.add(t)
  m -=pre.pop()
  m += d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

Solution (Python):
from collections import defaultdict
from heapq import heappush, heappop
N,K,*L = map(int, open(0).read().split())
dic = defaultdict(list)
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  dic[t].append(d)
  heappush(ls,(-d,t))
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = heappop(ls)
  ans -= d
  if t not in S:
    num += 1
    S.add(t)
  else:
    heappush(pre,-d)

ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = heappop(ls)
  if t in S:
    continue
  S.add(t)
  m -= heappop(pre)
  m += -d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

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

NETA=set()
duplication=[]
heapq.heapify(duplication)
ans=0
for t,d in TD[:K]:
    if t in NETA:
        heapq.heappush(duplication,d)
    else:
        NETA.add(t)
    ans+=d

s=[-float("inf")]*(K+1)
cnt=len(NETA)
s[cnt]=ans
for t,d in TD[K:]:
    if t in NETA:
        continue
    else:
        if duplication:
            x=heapq.heappop(duplication)
            ans+=d-x
            NETA.add(t)
            cnt+=1
            s[cnt]=ans

for i in range(K+1):
    s[i]+=i*i
print(max(s))

Solution (Python):
import heapq
n,k=map(int,input().split())
td=[[0]*2 for _ in range(n)]
for i in range(n):
    t,d=map(int,input().split())
    td[i][0],td[i][1]=d,t
td.sort(reverse=True)
di=set()
sa=[]
heapq.heapify(sa)
ans=0
for d,t in td[:k]:
    if t in di:
        heapq.heappush(sa, d)
    else:
        di.add(t)
    ans+=d

s=[-10**18]*(k+1)
cnt=len(di)
s[cnt]=ans
for d,t in td[k:]:
    if t in di:
        continue
    else:
        if sa:
            x=heapq.heappop(sa)
            ans+=d-x
            di.add(t)
            cnt+=1
            s[cnt]=ans

for i in range(k+1):
    s[i]+=i*i

print(max(s))

Solution (Python):
from operator import itemgetter
n,k=map(int,input().split())
td=[list(map(int,input().split())) for _ in range(n)]
td.sort(key=itemgetter(1),reverse=True)
ans=0
kind=0
flg=[False]*n
sing,comp=[],[]
for i in range(n):
    t,d=td[i]
    t-=1
    if i<k:
        ans+=d
        if flg[t]:
            comp.append(d)
        else:
            flg[t]=True
            kind+=1
    else:
        if not flg[t]:
            sing.append(d)
            flg[t]=True
comp=comp[::-1]
cnt=ans
ans+=kind**2


if len(comp)>len(sing):
    for i in range(len(sing)):
        dc,ds=comp[i],sing[i]
        cnt+=ds-dc
        kind+=1
        ans=max(ans,cnt+kind**2)
else:
    for i in range(len(comp)):
        dc,ds=comp[i],sing[i]
        cnt+=ds-dc
        kind+=1
        ans=max(ans,cnt+kind**2)
print(ans)

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)
cand = []
memo = [0 for i in range(n+1)]

ans = 0
duplicate = []
unique = []
num = 0
i = 0

for t, d in td:
    i += 1
    if i <= k:
        ans += d
        if memo[t]:
            duplicate.append(d)
        else:
            memo[t] = 1
            num += 1
    else:
        if not memo[t]:
            unique.append(d)
            memo[t] = 1

tmp = ans
ans += num * num



for d, u in zip(reversed(duplicate), unique):
    tmp += u - d
    num += 1
    ans = max(ans, tmp + num * num)

print(ans)

Solution (Python):
def main():
    n, k = map(int, input().split())
    sushi = [list(map(int, input().split())) for _ in range(n)]
    sushi.sort(key=lambda x: x[1], reverse=True)
    ans = 0
    t = [0]*(n+1)
    tmp = 0
    cnt = 0
    for i in range(k):
        ans += sushi[i][1]
        if t[sushi[i][0]] == 0:
            cnt += 1
        tmp += sushi[i][1]
        t[sushi[i][0]] += 1
    ans = tmp + cnt*cnt
    p = k
    for i in range(k, n):
        f = False
        if t[sushi[i][0]] == 0:
            f = True
            for j in reversed(range(p)):
                if 1 < t[sushi[j][0]]:
                    t[sushi[j][0]] -= 1
                    t[sushi[i][0]] += 1
                    tmp -= sushi[j][1]
                    tmp += sushi[i][1]
                    p = j
                    cnt += 1
                    ans = max(ans, tmp+cnt*cnt)
                    f = False
                    break
        if f:
            break
    print(ans)


if __name__ == "__main__":
    main()

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

N, K = map(int, input().split())
td = [tuple(map(int, input().split())) for _ in range(N)]
td.sort(key=lambda t: t[1], reverse=True)
d = defaultdict(list)
now = 0

for ti, di in td[:K]:
    d[ti].append(di)
    now += di

remove = []

for k in d.keys():
    d[k].sort(reverse=True)
    
    for i in range(1, len(d[k])):
        heappush(remove, d[k][i])
    
s = set(list(d.keys()))
add = []

for ti, di in td:
    if ti in s:
        continue
    
    s.add(ti)
    heappush(add, -di)

x = len(d)
ans = now+x**2

while len(remove)>0 and len(add)>0:
    now -= heappop(remove)
    now -= heappop(add)
    x += 1
    ans = max(ans, now+x**2)

print(ans)

Solution (Python):





inf = 10**9+7
mod = 10**9+7
n, k = map(int, input().split())
g_val = [list(map(int, input().split())) for i in range(n)]
netalist = [False for i in range(n+1)]
netanum = [0 for i in range(n+1)]
g_val.sort(key=lambda x: x[1])
used_list = []
ans = [0 for i in range(k+1)] 
preans = 0
for j in range(k):
    if g_val:
        g,v = g_val.pop()
        netalist[g] = True
        netanum[g]+=1
        used_list.append((v,g))
        preans+=v

kind = sum(list(map(int, netalist)))
ans[kind] = preans+(kind*kind)

import heapq 
heapq.heapify(used_list)   


while used_list:
    v,g = heapq.heappop(used_list) 
    if netanum[g] > 1:
        while g_val:
            gnew,vnew = g_val.pop()
            if netanum[gnew] == 0:
                ans[kind+1] = ans[kind]-v-(kind*kind)+vnew+((kind+1)*(kind+1))
                kind+=1
                netanum[g] -= 1
                netanum[gnew] += 1
                break
    if not g_val:
        break
print(max(ans))

Solution (Python):
n,k = map(int,input().split())
s = set()
l = []
ln = [[] for i in range(n)]
for i in range(n):
    t,m = map(int,input().split())
    l.append([m,t-1])
    ln[t-1].append(m)
l.sort(reverse = True)
cand = []

ans = 0
mod = []
for i in range(k):
    ans += l[i][0]
    if l[i][1] in s:
        mod.append(l[i][0])
    else:
        s.add(l[i][1])
x = len(s)
now = ans + x**2

for i in range(n):
    
    if i in s or ln[i] == []:
        continue
    cand.append(max(ln[i]))

cand.sort(reverse = True)

for i in range(min(len(mod),len(cand))):
    ans += cand[i]-mod[-1-i]
    x += 1
    now = max(now,ans+x**2)
print(now)

Solution (Python):

import sys
from heapq import heapify, heappop, heappush

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)]
heap = []
kind = 0
happy = 0
used = set()
TD.sort(key=lambda x: x[1], reverse=True)
for t, d in TD[:K]:
    if t not in used:
        used.add(t)
        kind += 1
    else:
        heappush(heap, d)
    happy += d

cand = [happy+kind**2]
for t, d in TD[K:]:
    if not heap:
        break
    if t in used:
        continue
    h = heappop(heap)
    happy += (d-h)
    kind += 1
    used.add(t)
    cand.append(happy+kind**2)

answer = max(cand)
print(answer)

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

dtlist=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tset_all.add(t)
dtlist.sort()


dtlist_Krev=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  d,t=dtlist.pop()
  dtlist_Krev.append((d,t))
  dsum+=d
  tdic[t]+=1


t0=len(tdic)
max_answer=dsum+(t0**2)


for i in range(t0+1,min(K,len(tset_all))+1):
  while(True):    
    d1,t1=dtlist_Krev.pop()
    if tdic[t1]>1:
      tdic[t1]-=1
      dsum-=d1
      break
      
  while(True):        
    d2,t2=dtlist.pop()
    if tdic[t2]==0:
      tdic[t2]=1
      dsum+=d2      
      break
          
  answer_i=dsum+(i**2)
  
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

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

dtlist=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tset_all.add(t)
dtlist.sort()


dtlist_Krev=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  d,t=dtlist.pop()
  dtlist_Krev.append((d,t))
  dsum+=d
  tdic[t]+=1
dtlist_Krev.sort(reverse=True)

t0=len(tdic)
max_answer=dsum+(t0**2)


for i in range(t0+1,min(K,len(tset_all))+1):
  loop_flg=True
  while(loop_flg):    
    d,t=dtlist_Krev.pop()
    if tdic[t]>1:
      tdic[t]-=1
      dsum-=d
      
      while(True):        
        d2,t2=dtlist.pop()
        if tdic[t2]==0:
          tdic[t2]=1
          dsum+=d2
          loop_flg=False
          break
          
  answer_i=dsum+(i**2)
  
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

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

hq=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  heapq.heappush(hq,(-d,t))
  tset_all.add(t)


hq_K=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  md,t=heapq.heappop(hq)
  heapq.heappush(hq_K,((-md,t)))
  dsum-=md
  tdic[t]+=1  

t0=len(tdic)
max_answer=dsum+(t0**2)


for i in range(t0+1,min(K,len(tset_all))+1):
  loop_flg=True
  while(loop_flg):
    d,t=heapq.heappop(hq_K)
    if tdic[t]>1:
      tdic[t]-=1
      dsum-=d
      
      while(True):
        md2,t2=heapq.heappop(hq)
        if tdic[t2]==0:
          tdic[t2]=1
          dsum-=md2
          loop_flg=False
          break
          
  answer_i=dsum+(i**2)
  
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

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

n, k = map(int, input().split())
TD = [[0, 0] for _ in range(n)]
for i in range(n):
    t, d = map(int, input().split())
    TD[i][0], TD[i][1] = d, t

TD.sort(reverse=True)
S = set()
f = 0
h = []
import heapq
heapq.heapify(h)
for i in range(k):
    d, t = TD[i]
    if t in S:
        heapq.heappush(h, d)
    else:
        S.add(t)
    f += d

INF = 10**18
F = [-INF]*(k+1)
cur = len(S)
F[cur] = f

for i in range(k, n):
    d, t = TD[i]
    if t in S:
        continue
    else:
        if h:
            x = heapq.heappop(h)
            f += d-x
            S.add(t)
            cur += 1
            F[cur] = f

for i in range(0, k+1):
    F[i] += i*i

print(max(F))

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

n, k = map(int, input().split())
TD = [[0, 0] for _ in range(n)]
for i in range(n):
    t, d = map(int, input().split())
    TD[i][0], TD[i][1] = d, t

TD.sort(reverse=True)
S = set()
f = 0
h = []
import heapq
heapq.heapify(h)
for i in range(k):
    d, t = TD[i]
    if t in S:
        heapq.heappush(h, d)
    else:
        S.add(t)
    f += d

INF = 10**18
F = [-INF]*(k+1)
cur = len(S)
F[cur] = f

for i in range(k, n):
    d, t = TD[i]
    if t in S:
        continue
    else:
        if h:
            x = heapq.heappop(h)
            f += d-x
            S.add(t)
            cur += 1
            F[cur] = f

for i in range(0, k+1):
    F[i] += i*i

print(max(F))

Solution (Python):


from collections import defaultdict

def main():
    n, k = map(int, input().split())
    td = []
    for i in range(n):
        ti, di = map(int, input().split())
        td.append((ti, di))
    td = list(reversed(sorted(td)))
    tdd = defaultdict(list)
    for t, d in td:
        tdd[t].append(d)

    a = []; b = []
    for _, v in tdd.items():
        a.append(v[0])
        b += v[1:]

    a = list(reversed(sorted(a)))
    b = list(reversed(sorted(b)))
    for i in range(1, len(a)):
        a[i] += a[i - 1]
    for i in range(len(a)):
        a[i] += (i + 1) ** 2
    for i in range(1, len(b)):
        b[i] += b[i - 1]

    a = [0] + a
    b = [0] + b

    res = 0
    for i in range(1, k + 1):
        if i >= len(a) or (k - i) >= len(b):
            continue
        ts = a[i] + b[k - i]
        res = max(res, ts)

    print(res)

if __name__ == "__main__":
    main()

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):
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)
a = 0
ans = []
kind = [-1]*n
c = 0
min_p = []
for i in range(n):
    eat = False
    if i<k:
        a += td[i][1]
        eat = True
    else:
        if min_p:
            if kind[td[i][0]-1]<0:
                a -= (min_p.pop()-td[i][1])
                a += (c+1)**2-c**2
                ans.append(a)
                eat = True
    if eat:
        if kind[td[i][0]-1]<0:
            kind[td[i][0]-1]=1
            c+=1
        else:min_p.append(td[i][1])
    if i==k-1:
        a += c**2
        ans.append(a)
print(max(ans))

Solution (Python):
n,k=map(int,input().split())
td=[]
for _ in range(n):
  td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
td_core=[]
td_sel=[]
td_sub=[]
cat=[0]*(n+1)
for i in range(n):
  if i < k and cat[td[i][0]]==0:
    td_core.append(td[i])
    cat[td[i][0]]=1
  elif i < k:
    td_sel.append(td[i])
  elif cat[td[i][0]]==0:
    td_sub.append(td[i])
    cat[td[i][0]]=1




kiso=sum([x[1] for x in td_core])+sum([x[1] for x in td_sel])
bns=len(td_core)
ans=kiso+bns**2



td_sel.sort(key=lambda x:x[1])
for i in range(min(len(td_sel),len(td_sub))):
  kiso-=td_sel[i][1]
  kiso+=td_sub[i][1]
  bns+=1
  ans=max(ans,kiso+bns**2)
print(ans)

Solution (Python):

from collections import defaultdict, deque
from heapq import heappush, heappop
from itertools import permutations, accumulate
import sys
import math
import bisect
def LI(): return [int(x) for x in sys.stdin.buffer.readline().split()]
def I(): return int(sys.stdin.buffer.readline())
def LS():return [list(x) for x in sys.stdin.readline().split()]
def S():
    res = list(sys.stdin.readline())
    if res[-1] == "\n":
        return res[:-1]
    return res
def IR(n):
    return [I() for i in range(n)]
def LIR(n):
    return [LI() for i in range(n)]
def SR(n):
    return [S() for i in range(n)]
def LSR(n):
    return [LS() for i in range(n)]

sys.setrecursionlimit(1000000)
mod = 1000000007

def solve():
    n,k = LI()
    g = LIR(n)
    l = list(set(([i for (i,j) in g])))
    g.sort(key=lambda x:-x[1])
    l.sort()
    q = []
    q2 = []
    f = [0]*len(l)
    ans = 0
    res = k
    for i in range(n):
        t,d = g[i]
        g[i][0] = t = bisect.bisect_left(l,t)
        if res:
            f[t] += 1
            ans += d
            res -= 1
            if f[t] > 1:
                heappush(q,(d,t))
        else:
            heappush(q2,(-d,t))
    x = sum([i >= 1 for i in f])
    ans += x**2
    s = ans
    while q2 and q:
        d,t = heappop(q2)
        d *= -1
        if f[t]:
            continue
        else:
            nd,nt = heappop(q)
            s -= nd
            s += 2*x+1
            s += d
            x += 1
            f[t] = 1
            f[nt] -= 1
            if ans < s:
                ans = s
    print(ans)
    return


if __name__ == "__main__":
    solve()

Solution (Python):
import sys
from heapq import heapify, heappop, heappush

def solve():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    S = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        S[i] = (-d, t)
    heapify(S)

    base_point = 0
    appeared_type = set()
    type_Num = 0
    duplicated = []
    heapify(duplicated)
    for i in range(K):
        d, t = heappop(S)
        if t in appeared_type:
            base_point -= d
            heappush(duplicated, -d)
        else:
            base_point -= d
            appeared_type |= {t}
            type_Num += 1

    Ans = base_point + type_Num * type_Num
    while S:
        d, t = heappop(S)
        if t in appeared_type: continue
        else:
            if not duplicated: continue
            else:
                remove = heappop(duplicated)
                base_point -= d + remove
                type_Num += 1
                appeared_type |= {t}
                Ans = max(Ans, base_point + type_Num * type_Num)
    print(Ans)

    return 0

if __name__ == "__main__":
    solve()

Solution (Python):
n,k=map(int,input().split())
td=[list(map(int,input().split()))[::-1]for _ in range(n)]
td.sort(reverse=1)
s=set()
ss=set()
a=[]
b=[]
c=[]
for i in range(n):
  if i<k:
    d,t=td[i]
    if t in s:
      b.append(d)
    else:
      a.append(d)
      s.add(t)
      ss.add(t)
  else:
    d,t=td[i]
    if t not in ss:
      c.append(d)
      ss.add(t)
c=c[::-1]
ls=len(s)
su=sum(a)+sum(b)
ans=ls*ls+su
for i in range(min(len(b),len(c))):
  ls+=1
  su-=b.pop()
  su+=c.pop()
  ans=max(ans,ls*ls+su)
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])

ans=0
kl=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kl:
        kl[td[i][0]]+=1
    else:
        kl[td[i][0]]=1
l=len(kl)
ans+=(l**2)
ans_=ans
now=k-1
for i in range(k,n):
    if td[i][0] not in kl:
        
        while now>=0:
            if kl[td[now][0]]>1:
                mi=td[now]
                kl[mi[0]]-=1
                now-=1
                break
            now-=1
        if now==-1:
            break
        else:
            ans=ans+2*l+1-mi[1]+td[i][1]
            kl[td[i][0]]=1
            ans_=max(ans,ans_)
            l+=1

print(ans_)

Solution (Python):
def Int():
    return int(input())
def Ints():
    return map(int,input().split())
def IntList():
    return list(Ints())
def IntMat(N):
    return [IntList() for i in range(N)]

import sys
sys.setrecursionlimit(4100000)
input = sys.stdin.readline

N,K = Ints()
ALL = []
NetaMax = 0
MaxNeta = {}

for i in range(N):
    t,d = Ints()
    ALL.append([d,t])
    if not t in MaxNeta:
        NetaMax += 1
        MaxNeta[t] = d
    else:
        MaxNeta[t] = max(MaxNeta[t], d)
    
ALL.sort()
ALL = ALL[::-1]
NetaBest = []

for i in MaxNeta:
    NetaBest.append(MaxNeta[i])
    
NetaBest.sort()
NetaBest = NetaBest[::-1]

ans = 0
tmpans = 0
count = 0
tmpD = {}
Second = []

for i in range(K):
    d, t = ALL[i]
    if not t in tmpD:
        tmpD[t] = 1
        count += 1
    else:
        Second.append(d)
    tmpans += d

Second.sort()

ans = tmpans + count**2


for i in range(count+1, min(K, NetaMax)+1):
    
    tmpans -= Second[i - (count+1)]
    tmpans += NetaBest[i-1]
    
    ans = max(ans, tmpans + i**2)
    
print(ans)

Solution (Python):
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []
for key in sushi:
    sushi[key].sort(reverse=True)
    first.append(sushi[key][0])
    second += sushi[key][1:]

first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))


ans = 0
for i in range(1, K+1):
    if K - i <= len(second_cum)-1 and i <= len(first_cum)-1:
        ans = max(ans, first_cum[i] + second_cum[K-i] + i*i)

print(ans)

Solution (Python):
import heapq as hq

n,k=map(int,input().split())
A=[]
S0=[] 
S1=[] 
for i in range(n):
    t,d=map(int,input().split())
    A.append([t,d])
A=list(reversed(sorted(A)))

now=0
for i in range(n):
    t=A[i][0]
    d=A[i][1]
    if now==t:
        S0.append(-d)
    else:
        S1.append(-d)
        now=t
hq.heapify(S0)
hq.heapify(S1)
C=[] 

x=0 
y=0 
for i in range(k):
    if len(S0)==0:
        x+=-hq.heappop(S1)
        y+=1
        continue
    if len(S1)==0:
        s0=-hq.heappop(S0)
        x+=s0
        hq.heappush(C,s0)
        continue
    if S1[0]<=S0[0]:
        x+=-hq.heappop(S1)
        y+=1
    else:
        s0=-hq.heappop(S0)
        x+=s0
        hq.heappush(C,s0)
ans=x+y**2

for i in range(min(len(S1),len(C))):
    c=hq.heappop(C)
    s1=-hq.heappop(S1)
    x+=s1-c
    y+=1
    ans=max(x+y**2,ans)
print(ans)

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

Solution (Python):
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

for i in range(N):
    a,b = map(int, input().split())
    sushi[a].append(b)

first = []
second = []

for nums in sushi.values():
    nums.sort(reverse=True)
    for i, j in enumerate(nums):
        if i == 0:
            first.append(j)
        else:
            second.append(j)
first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))

ans = 0

for i in range(K, 0, -1):
    neta = min(i, len(first))
    tmp = neta * neta + first_cum[neta]

    if K - neta > len(second):
        break
    
    tmp += second_cum[K - neta]

    ans = max(ans, tmp)

print(ans)

Solution (Python):
n,k= map(int,input().split())
td=[]
for _ in range(n):
    td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
NEW=[1]*(n+1)
x=0
y=0
TOP=[]
SUB=[]
A=[]
for tdi in td[:k]:
    if NEW[tdi[0]]==1:
        NEW[tdi[0]]=0
        TOP.append(tdi[1])
        y+=1
    else:
        A.append(tdi[1])
    x+=tdi[1]
    
for tdi in td[k:]:
    if NEW[tdi[0]]==1:
        SUB.append(tdi[1])
        NEW[tdi[0]]=0
ans=x+y**2
for i in range(min(k-1,len(A),len(SUB))):
    x -= A.pop()
    x += SUB[i]
    y += 1
    ans = max(ans,x+y**2)
print(ans)

Solution (Python):
from heapq import heappush,heappop
n,k = map(int,input().split())
kind_value = [[] for _ in range((n+1))]
for i 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):
from heapq import heappush,heappop
n,k = map(int,input().split())
kind_value = [[] for _ in range((n+1))]
for i 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):
def main():
    import sys
    input = sys.stdin.readline
    sys.setrecursionlimit(10**7)
    from collections import Counter, deque
    
    from itertools import combinations, permutations, accumulate, groupby
    
    from bisect import bisect_left,bisect_right
    from heapq import heapify, heappop, heappush
    from math import floor, ceil
    

    
    

    n,k = map(int, input().split())
    td = [list(map(int, input().split())) for _ in range(n)]
    td.sort(reverse=True, key=lambda a: a[1])
    delicious = 0
    syurui = set()
    suteru = []
    for i in range(k):
        t, d = td[i]
        delicious += d
        if  t in syurui:
            suteru.append(d)
        else:
            syurui.add(t)
    
    suteru.sort()
    neta = len(syurui)
    res = delicious + neta**2
    for i in range(k, n):
        t, d = td[i]
        if t in syurui:
            continue
        if len(suteru)==0:
            print(res)
            exit()
        syurui.add(t)
        x = suteru.pop(0)
        delicious += d - x
        neta += 1
        res = max(res, delicious+neta**2)
    print(res)

if __name__ == '__main__':
    main()

Solution (Python):
def main():
    import sys
    input = sys.stdin.readline
    sys.setrecursionlimit(10**7)
    from collections import Counter, deque
    
    from itertools import combinations, permutations, accumulate, groupby
    
    from bisect import bisect_left,bisect_right
    from heapq import heapify, heappop, heappush
    from math import floor, ceil
    

    
    

    n,k = map(int, input().split())
    td = [list(map(int, input().split())) for _ in range(n)]
    td.sort(reverse=True, key=lambda a: a[1])
    delicious = 0
    syurui = set()
    suteru = []
    for i in range(k):
        t, d = td[i]
        delicious += d
        if  t in syurui:
            suteru.append(d)
        else:
            syurui.add(t)
    heapify(suteru)
    neta = len(syurui)
    res = delicious + neta**2
    for i in range(k, n):
        t, d = td[i]
        if t in syurui:
            continue
        if len(suteru)==0:
            print(res)
            exit()
        syurui.add(t)
        delicious += d - heappop(suteru)
        neta += 1
        res = max(res, delicious+neta**2)
    print(res)

if __name__ == '__main__':
    main()