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

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


ans = sum(selected0) + sum(selected1) + various ** 2

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

print(ans)

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

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)

base = []; other = []
val = set()
yum = 0
for t, d in td[:k]:
    yum += d
    if t in val:
        other.append(d)
    val.add(t)
other.sort(reverse=True)

kouho = []
val_tmp = val.copy()
for t, d in td[k:]:
    if t not in val_tmp:
        kouho.append(d)
        val_tmp.add(t)
kouho.sort()

v = len(val)
ans = yum + (v*v)


while len(other) > 0 and len(kouho) > 0:
    yum -= other.pop()
    yum += kouho.pop()
    v += 1
    ans = max(ans, yum+(v*v))
print(ans)

Solution (Python):

import heapq

n,k,*td=map(int,open(0).read().split())
td = [[t,d] for t,d in zip(td[::2],td[1::2])]
td = sorted(td,key = lambda x:x[1],reverse=True)

score = [0] * k
taste_sum = 0
extra_taste = []
eat_num = 0
eat_kind = set()

for t,d in td:
    if t in eat_kind:
        if eat_num == k:
            continue
        else:
            heapq.heappush(extra_taste, d)
            taste_sum += d
            eat_num += 1
    else:
        eat_kind.add(t)
        taste_sum += d
        eat_num += 1
        if eat_num > k:
            eat_num -= 1
            taste_sum -= heapq.heappop(extra_taste)
    
    if eat_num == k:
        score[len(eat_kind)-1] = len(eat_kind)**2 + taste_sum
            
    if len(eat_kind) == k:
        break
print(max(score))

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

N, K = map(int, input().split())
TD = []

sushi = []
kinds = []


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

sushi.sort(key=lambda x: x[0], reverse=True)
sushidict = defaultdict(int)
kind_all = len(set(kinds))

cnt = 0
sushiq = []

for i in range(K):
  d, t = sushi[i]
  heappush(sushiq, (d, t))
  sushidict[t] += 1
  cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(K, N):
  d, t = sushi[i]

  if sushidict[t] > 0:
    continue

  d_min, t_min = heappop(sushiq)
  while sushidict[t_min] == 1 and sushiq:
    d_min, t_min = heappop(sushiq)

  if not sushiq:
    break

  sushidict[t_min] -= 1
  sushidict[t] += 1
  n_kind += 1
  cnt = cnt - d_min + d

  res.append(cnt+n_kind ** 2)
  if n_kind == kind_all:
    break

print(max(res))

Solution (Python):

import sys
import math
import os
import itertools
import string
import heapq
import _collections
from collections import Counter
from collections import defaultdict
from functools import lru_cache
import bisect
import re
import queue


class Scanner():
    @staticmethod
    def int():
        return int(sys.stdin.readline().rstrip())

    @staticmethod
    def string():
        return sys.stdin.readline().rstrip()

    @staticmethod
    def map_int():
        return [int(x) for x in Scanner.string().split()]

    @staticmethod
    def string_list(n):
        return [input() for i in range(n)]

    @staticmethod
    def int_list_list(n):
        return [Scanner.map_int() for i in range(n)]

    @staticmethod
    def int_cols_list(n):
        return [int(input()) for i in range(n)]


class Math():
    @staticmethod
    def gcd(a, b):
        if b == 0:
            return a
        return Math.gcd(b, a % b)

    @staticmethod
    def lcm(a, b):
        return (a * b) // Math.gcd(a, b)

    @staticmethod
    def roundUp(a, b):
        return -(-a // b)

    @staticmethod
    def toUpperMultiple(a, x):
        return Math.roundUp(a, x) * x

    @staticmethod
    def toLowerMultiple(a, x):
        return (a // x) * x

    @staticmethod
    def nearPow2(n):
        if n <= 0:
            return 0
        if n & (n - 1) == 0:
            return n
        ret = 1
        while(n > 0):
            ret <<= 1
            n >>= 1
        return ret

    @staticmethod
    def sign(n):
        if n == 0:
            return 0
        if n < 0:
            return -1
        return 1

    @staticmethod
    def isPrime(n):
        if n < 2:
            return False
        if n == 2:
            return True
        if n % 2 == 0:
            return False
        d = int(n ** 0.5) + 1
        for i in range(3, d + 1, 2):
            if n % i == 0:
                return False
        return True


MOD = int(1e09) + 7
INF = int(1e15)

T_MAX = int(1e05) + 1


def main():
    
    N, K = Scanner.map_int()
    T = defaultdict(list)
    for _ in range(N):
        t, d = Scanner.map_int()
        t -= 1
        T[t].append(d)
    for v in T.values():
        v.sort(reverse=True)
    M = []
    for v in T.values():
        M.append((v[0], 1))
        for i in range(1, len(v)):
            M.append((v[i], 0))
    M.sort(reverse=True)

    q = []
    x = 0
    y = 0
    for i in range(K):
        x += M[i][0]
        y += M[i][1]
        if M[i][1] == 0:
            heapq.heappush(q, M[i][0])
    ans = x + y * y
    for i in range(K, N):
        if q == []:
            break
        if M[i][1] == 1:
            x -= heapq.heappop(q)
            x += M[i][0]
            y += 1
            ans = max(ans, x + y * y)
    print(ans)


if __name__ == "__main__":
    main()

Solution (Python):
import sys
from collections import *
import heapq
import math
import bisect
from itertools import permutations,accumulate,combinations,product
from fractions import gcd
def input():
    return sys.stdin.readline()[:-1]
mod=10**9+7

n,k=map(int,input().split())
td=[list(map(int,input().split())) for i in range(n)]
td.sort(key=lambda x: -x[1])

dic={}
minlst=[]
maxlst=[]
ans=0
f=0
ln=0
for i in range(n):
    t,d=td[i]
    if i<k:
        f+=d
        if not t in dic:
            dic[t]=d
        else:
            minlst.append(d)
        ln=len(dic)
    else:
        if not t in dic:
            maxlst.append(-d)
            dic[t]=d
        else:
            pass







ans=max(ans,f+pow(ln,2))
heapq.heapify(minlst)
heapq.heapify(maxlst)
while minlst and maxlst:
    f-=heapq.heappop(minlst)
    f-=heapq.heappop(maxlst)
    ln+=1
    ans=max(ans,f+pow(ln,2))
print(ans)

Solution (Python):
Num,K=map(int,input().split())
L=[[]for i in range(Num+1)]
for i in range(Num):
    t,d=map(int,input().split())
    L[t].append(d)
for i in range(1,Num+1):
    L[i].sort(reverse=True)

M=[]
m=[]
for i in range(1,Num+1):
    if len(L[i])>=1:
        M.append(L[i][0])
        if len(L[i])>=2:
            for j in range(1,len(L[i])):
                m.append(L[i][j])
M.sort(reverse=True)
m.sort(reverse=True)
MM=[0]
for i in range(len(M)):
    MM.append(MM[-1]+M[i])
mm=[0]
for i in range(len(m)):
    mm.append(mm[-1]+m[i])


ANS=0

def cmb(n, r, mod):
    if ( r<0 or r>n ):
        return 0
    r = min(r, n-r)
    return g1[n] * g2[r] * g2[n-r] % mod

mod = 10**9+7 
N = 10**5+10
g1 = [1, 1] 
g2 = [1, 1] 
inverse = [0, 1] 

for i in range( 2, N + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
    inverse.append( ( -inverse[mod % i] * (mod//i) ) % mod )
    g2.append( (g2[-1] * inverse[-1]) % mod )

for i in range(1,min(K,len(M))+1):
    ANS=max(i*i+MM[i]+mm[min(len(mm)-1,K-i)],ANS)
    
print(ANS)

Solution (Python):
N,K = map(int,input().split())
ans = [0 for i in range(N+1)]
food = {}
food1 = []
for i in range(1,N+1):
    food[i] = 0
for i in range(N):
    t,d = map(int,input().split())
    food1.append((d,t))
food1.sort()
syurui = 0
sum = 0
import heapq
x = []
heapq.heapify(x)
for i in range(1,N+1):
    d,t = food1[-i]
    if food[t] == 0:
        if i > K:
            if len(x) == 0:
                break
        sum += d
        food[t] = 1
        syurui += 1
        if i > K:
            y = heapq.heappop(x)
            sum -= y
    else:
        if i <= K:
            sum += d
            heapq.heappush(x,d)
    if syurui == K:
        ans[K] = max(sum + K**2,ans[K])
        break
    else:
        ans[syurui] = max(sum + syurui**2,ans[syurui])
print(max(ans))

Solution (Python):



import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N+1)]
x=0
ans=0

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    dt[i][1]-=1
    ans+=dt[i][0]
    if gr[dt[i][1]]==[]:
        x+=1
    gr[dt[i][1]].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    
inf=10**19
heapq.heappush(single,[inf,N])
heapq.heappush(multi,[inf,N])
    
    
    
for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        tt=b[1]
        num_i[tt]+=1
        if num_i[tt]==len(gr[tt])-1:
            heapq.heappush(single,gr[tt][-1])
        else:
            heapq.heappush(multi,gr[tt][num_i[tt]])
                                       
print(ans)

Solution (Python):











import sys
sys.setrecursionlimit(10**6)
input=sys.stdin.readline
from math import floor,ceil,sqrt,factorial,hypot,log 
from heapq import heappop, heappush, heappushpop
from collections import Counter,defaultdict,deque
from itertools import accumulate,permutations,combinations,product,combinations_with_replacement
from bisect import bisect_left,bisect_right
from copy import deepcopy
from fractions import gcd
from random import randint
inf=float('inf')
mod = 10**9+7
def pprint(*A): 
    for a in A:     print(*a,sep='\n')
def INT_(n): return int(n)-1
def MI(): return map(int,input().split())
def MF(): return map(float, input().split())
def MI_(): return map(INT_,input().split())
def LI(): return list(MI())
def LI_(): return [int(x) - 1 for x in input().split()]
def LF(): return list(MF())
def LIN(n:int): return [I() for _ in range(n)]
def LLIN(n: int): return [LI() for _ in range(n)]
def LLIN_(n: int): return [LI_() for _ in range(n)]
def LLI(): return [list(map(int, l.split() )) for l in input()]
def I(): return int(input())
def F(): return float(input())
def ST(): return input().replace('\n', '')
def main():
    N,K = MI()
    TD = LLIN(N)
    TD.sort(key = lambda x:x[1], reverse=True)

    neta = set()
    double = []
    ans = 0
    S = 0
    for i,(t,d) in enumerate(TD):
        if i < K:
            if t in neta:
                double.append(d)
            else:
                neta.add(t)
            S += d
        else:
            if t not in neta and double:
                S -= double.pop()
                S += d
                neta.add(t)
        ans = max(ans, S + len(neta)**2)
    print(ans)
if __name__ == '__main__':
    main()

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


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)

base = 0
ate = Counter()
q = []
for i in range(K):
    t, d = sushi[i]
    base += d
    heappush(q, (d, t))
    ate[t] += 1

var = len(list(ate.keys()))
candidate = [base + var ** 2]
for i in range(K, N):
    t, d = sushi[i]
    if ate[t] == 0:
        while len(q):
            dx, tx = heappop(q)
            if ate[tx] > 1:
                break
        else:
            break
        ate[t] += 1
        ate[tx] -= 1
        base += d - dx
        var += 1
        candidate.append(base + var ** 2)

print(max(candidate))

Solution (Python):
from operator import itemgetter
from itertools import accumulate


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

sushi.sort(key=itemgetter(1))
sushi.reverse()



a = []


b = []


in_a = [0]*(n+1)

for i in range(n):
    t,d = sushi[i]
    if in_a[t]==0:
        in_a[t]=1
        a.append(d)
    else:
        b.append(d)

ans = 0

a = list(accumulate([0]+a))
b = list(accumulate([0]+b))


for i in range(len(a)):
    if i<=k:
        if 0<=k-i<len(b):
            ans = max(a[i]+b[k-i]+i**2,ans)
print(ans)

Solution (Python):
n,k = map(int,input().split())
l = [tuple(map(int, input().split())) for _ in range(n)]
l.sort(key=lambda x:-x[1])
dic = dict()
tr = list()
curd = 0
curt = 0
for t,d in l[:k]:
  curd += d
  if t not in dic:
    curt += 1
    dic[t] = d
  elif dic[t] < d:
    tr.append(dic[t])
    dic[t] = d
  else:
    tr.append(d)
ans = curt**2 + curd
tr.sort(reverse=True)
for t,d in l[k:]:
  if t in dic: continue
  if not tr: break
  curt += 1
  dic[t] = d
  curd += d - tr.pop()
  cur = curt**2 + curd
  if ans < cur: ans = cur
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[1], reverse=True)  
kind = set()
stack = []
ans = 0
for t, d in sushi[:K]:
    if t in kind:
        stack.append(d)
    else:
        kind.add(t)
    ans += d

ans += len(kind)**2
point = ans
for t, d in sushi[K:]:
    if len(stack) == 0:  
        break

    
    if t in kind:
        continue
    else:
        point -= stack.pop()
        point += d
        point += 2*len(kind) + 1  
        kind.add(t)
    ans = max(ans, point)
print(ans)

Solution (Python):
def examA():
    L = LI(); L.sort()
    ans = L[0]*L[1]//2
    print(ans)
    return

def examB():
    def fun(s):
        if s%2==0:
            return s//2
        else:
            return s*3+1
    S = I()
    d = defaultdict(bool)
    d[S] = True; cur = 1
    while(True):
        cur +=1
        S = fun(S)
        if d[S]:
            break
        d[S] = True
    print(cur)
    return

def examC():
    N = I()
    H = LI()
    cur = 0
    ans =0
    for i in range(N):
        if cur<H[i]:
            ans += H[i]-cur
        cur = H[i]
    print(ans)
    return

def examD():
    N, K = LI()
    TS = [LI() for _ in range(N)]
    choice = []
    keep = []
    d = defaultdict(int)
    TS.sort(key=lambda x:x[1],reverse=True)
    for i,(t,s) in enumerate(TS):
        if d[t]==0:
            if i < K:
                d[t] = -1
                choice.append(s)
            else:
                d[t] = s
        else:
            if i<K:
                keep.append(s)
    keep.sort()
    D = sorted(d.values(),reverse=True)

    cur = len(choice)
    now = cur**2 +sum(choice)+sum(keep)
    ans = copy.deepcopy(now)
    for i,k in enumerate(keep):
        if D[i]<0:
            break
        now += ((cur+i)*2+1) + (D[i]-k)
        ans = max(ans,now)

    print(ans)
    return

import sys,copy,bisect,itertools,heapq,math
from heapq import heappop,heappush,heapify
from collections import Counter,defaultdict,deque
def I(): return int(sys.stdin.readline())
def LI(): return list(map(int,sys.stdin.readline().split()))
def LSI(): return list(map(str,sys.stdin.readline().split()))
def LS(): return sys.stdin.readline().split()
def SI(): return sys.stdin.readline().strip()
global mod,mod2,inf,alphabet,ALPHABET, _ep
mod = 10**9 + 7
mod2 = 998244353
inf = 10**18
_ep = 1e-9
alphabet = [chr(ord('a') + i) for i in range(26)]
ALPHABET = [chr(ord('A') + i) for i in range(26)]

if __name__ == '__main__':
    examD()

Solution (Python):
n,k=[int(x) for x in input().split()]
A=[]
for i in range(n):
  a,b=[int(x) for x in input().split()]
  A.append([b,a,True])

A.sort(reverse=True)

typ=set()

for i in range(n):
  if A[i][1] in typ:
    A[i][2]=False
  else:
    typ.add(A[i][1])

num=0

for j in range(k):
  if A[j][2]==True:
    
    num+=1

ans=0
for i in range(k):
  ans+=A[i][0]
ans+=num**2



curr=ans
nxt = k
lst = k-1

while(True):

	while(nxt < n and A[nxt][2] == False):
		nxt += 1
	while(lst >= 0 and A[lst][2] == True):
		lst -= 1
	if not(nxt < n and lst >= 0):
		break

	curr += A[nxt][0] - A[lst][0]
	curr +=  (num+1) ** 2 - num**2


	num+=1
	nxt+= 1
	lst  -= 1

	ans=max(ans,curr)




print(ans)

Solution (Python):
def p_d():
    N, K = map(int, input().split())
    inf = 10 ** 18
    sushi = [[-inf] for _ in range(N)]
    max_kind = -1
    for _ in range(N):
        t, d = map(int, input().split())
        sushi[t - 1].append(d)
        if max_kind < t:
            max_kind = t
    for i in range(max_kind):
        sushi[i].sort(reverse=True)
    sushi.sort(key=lambda x: x[0], reverse=True)

    que = []
    res = 0
    for i in range(K):
        res += sushi[i][0]
        for j in sushi[i][1:]:
            que.append(j)
    que.sort()
    cur = res
    res += K ** 2
    for var in range(1, K)[::-1]:
        a = que.pop()
        b = sushi[var][0]
        cur += a - b
        res = max(res, cur + var ** 2)
    print(res)


if __name__ == '__main__':
    p_d()

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

n, k, *td = map(int, sys.stdin.read().split())
td = list(zip(*[iter(td)] * 2))

def main():
    td.sort(key=lambda x: x[1], reverse=True)
    delis = [[] for _ in range(n+1)]
    cnt = [0] * (n + 1)
    base = 0
    x = 0
    toppings = set()
    for t, d in td[:k]:
        toppings.add(t)
        if cnt[t] == 0:
            x += 1
        cnt[t] += 1
        heappush(delis[t], d)
        base += d
    score = base + x * x

    cand_to_remove = []
    for t in range(1, n+1):
        if cnt[t] >= 2:
            heappush(cand_to_remove, (heappop(delis[t]), t))
            cnt[t] -= 1

    for t, d in td[k:]:
        if not cand_to_remove:
            break
        if t in toppings:
            continue
        r_d, r_t = heappop(cand_to_remove)
        toppings.add(t)
        base -= r_d
        cnt[t] += 1
        heappush(delis[t], d)
        base += d
        x += 1
        if cnt[r_t] >= 2:
            heappush(cand_to_remove, (heappop(delis[r_t]), r_t))
            cnt[r_t] -= 1
        score = max(score, base + x * x)

    return score
        
if __name__ == '__main__':
    ans = main()
    print(ans)

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

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    
    y1.append(ds[i][-1])    
    ds[i].pop()
    y0.extend(ds[i])        
y0.sort(reverse=True)   
y1.sort(reverse=True)   

ans = 0

Y = min(k, len(y1))   
X = 0
for i in range(Y):
    X += y1[i]      
for i in range(k-Y):
    X += y0[i]      

while True:
    ans = max(ans, X+Y*Y)
    
    
    
    
    
    if (Y<=0 or k-Y>=len(y0)):   
        break
    X -= y1[Y-1]      
    X += y0[k-Y]  
    Y -= 1
print(ans)

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

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    
    y1.append(ds[i][-1])    
    ds[i].pop()
    y0.extend(ds[i])        
y0.sort(reverse=True)   
y1.sort(reverse=True)   

ans = 0
Y = max(0, k-len(y0))   
X = 0
for i in range(Y):
    X += y1[i]      
for i in range(k-Y):
    X += y0[i]      

while True:
    ans = max(ans, X+Y*Y)
    if (Y>=k or Y>=len(y1)):   
        break
    X += y1[Y]      
    X -= y0[k-Y-1]  
    Y += 1
print(ans)

Solution (Python):
n,k=map(int,input().split())
li=[list(map(int,input().split())) for i in range(n)]
gr=[0]*(n+1) 

from operator import itemgetter
li=sorted(li, key=itemgetter(1), reverse=True)

A=[] 

for i in range(k):
  if gr[li[i][0]]>=1:
    A.append(li[i][1])
  gr[li[i][0]]+=1

A.reverse()

B=[] 
for i in range(k,n):
  if gr[li[i][0]]==0:
    B.append(li[i][1])
  gr[li[i][0]]+=1


sp=set()
ans=0
for i in range(k):
  ans+=li[i][1]
  sp.add(li[i][0])
sp=len(sp)
ans+=sp*sp

koukan=[] 
zouka=0
for a,b in zip(A,B):
  zouka+=(sp+1)**2-sp**2-(a-b)
  
  koukan.append(zouka)
  sp+=1




if koukan:
  ans+=max(0,max(koukan))

print(ans)

Solution (Python):
import sys
from heapq import heappush,heappop,heapify
sys.setrecursionlimit(10**9)
INF=10**18
MOD=10**9+7
def input():
    return sys.stdin.readline().rstrip()

def main():
    N,K=map(int,input().split())
    sushies=[]
    for i in range(N):
        t,d=map(int,input().split())
        sushies.append((d,t-1))
    sushies.sort()
    b=[False]*N
    overlappedsushi=[]
    ans=0
    n=0
    for i in range(K):
        sushi=sushies.pop()
        if b[sushi[1]]:
            heappush(overlappedsushi,sushi)
            ans+=sushi[0]
        else:
            b[sushi[1]]=True
            ans+=sushi[0]
            n+=1
    ans+=n**2
    tmp=ans
    while overlappedsushi:
        worstsushi=heappop(overlappedsushi)
        tmp-=worstsushi[0]
        while sushies:
            bestsushi=sushies.pop()
            if not b[bestsushi[1]]:
                break
        else:
            break
        n+=1
        tmp+=bestsushi[0]
        tmp+=n*2-1
        b[bestsushi[1]]=True
        ans=max(ans,tmp)
    print(ans)

    

if __name__ == '__main__':
    main()

Solution (Python):
N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda p: -p[1])
neta = set()
oisi = netan = 0
amari = []
for t, d in TD[:K]:
    oisi += d
    if t not in neta:
        neta.add(t)
        netan += 1
    else:
        amari.append(d)
ans = oisi + netan * netan
for t, d in TD[K:]:
    if netan == K or not amari:
        break
    if t not in neta:
        neta.add(t)
        netan += 1
        oisi += d - amari.pop()
        ans = max(ans, oisi + netan * netan)
print(ans)

Solution (Python):
import os
import sys
from collections import deque
from operator import itemgetter

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7


N, K = list(map(int, sys.stdin.buffer.readline().split()))
TD = [list(map(int, sys.stdin.buffer.readline().split())) for _ in range(N)]

t_seen = [False] * (N + 1)
sushi_seen = [False] * N
TD.sort(key=itemgetter(1), reverse=True)

delicious = 0
que = deque()
for i, (t, d) in enumerate(TD):
    if len(que) >= K:
        break
    if not t_seen[t]:
        que.append(i)
        delicious += d
        t_seen[t] = True
        sushi_seen[i] = True
for i, (t, d) in enumerate(TD):
    if len(que) >= K:
        break
    if not sushi_seen[i]:
        que.appendleft(i)
        delicious += d
        sushi_seen[i] = True

idx = i
type_cnt = sum(t_seen)
ans = delicious + type_cnt ** 2
while idx < que[-1]:
    while idx < N and sushi_seen[idx]:
        idx += 1
    if idx >= N:
        break
    que.appendleft(idx)
    delicious -= TD[que.pop()][1]
    delicious += TD[idx][1]
    type_cnt -= 1
    score = delicious + type_cnt ** 2
    ans = max(ans, score)
    sushi_seen[idx] = True
print(ans)

Solution (Python):
import heapq
from operator import itemgetter

n,k = map(int, input().split())
td = sorted([list(map(int, input().split())) for i in range(n)],key=itemgetter(1),reverse=True)
neta = [0] * (n+1)
hq = []
kind = 0
ans = 0
for t,d in td[:k]:
  if neta[t]:
    hq.append(d)
  else:
    kind += 1
  neta[t] += 1
  ans += d
ans += kind ** 2
heapq.heapify(hq)
now = k
new_ans = ans
while hq and now < n:
  t,d = td[now][0],td[now][1]
  if not neta[t]:
    MIN = heapq.heappop(hq)
    new_ans += d - MIN + 2 * kind + 1
    ans = max(ans,new_ans)
    kind += 1
    neta[t] += 1
  now += 1	
print(ans)

Solution (Python):
import itertools

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

td = {}

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

top = []
rest = []

for k in td.keys():
    td[k].sort()
    top.append(td[k][-1])
    rest += td[k][:-1]

top.sort()
rest.sort()

cumtop = list(itertools.accumulate(reversed(top)))
cumrest = list(itertools.accumulate(reversed(rest)))

ret = 0

for i in range(min(len(top), K)):
    tmp = (i+1)*(i+1)
    tmp += cumtop[i]
    if i == K-1:
        pass
    else:
        if (K-i-2) < 0 or (K-i-2) > len(cumrest)-1:
            continue
        tmp += cumrest[(K-i-2)]
    ret = max(tmp, ret)

print(ret)

Solution (Python):
import heapq

N,K = map(int, input().split())
h = []
max_t = 0
for _ in range(N):
    t, d = map(int, input().split())
    heapq.heappush(h, (-d, t))
    max_t = max(max_t, t)

L = [0]*max_t
num = 0
h2 = []
for _ in range(K):
    temp = heapq.heappop(h)
    L[temp[1]-1] += 1
    num -= temp[0]
    heapq.heappush(h2, (-temp[0], temp[1]))
    
cnt = 0
for i in range(max_t):
    if L[i] > 0:
        cnt += 1

ans = num + cnt**2
while len(h) > 0 and len(h2) > 0:
    temp = heapq.heappop(h)
    if L[temp[1]-1] > 0:
        continue
    while L[h2[0][1]-1] <= 1:
        heapq.heappop(h2)
        if not h2:
            break
    if not h2:
        break
    temp2 = heapq.heappop(h2)    
    num -= temp2[0] + temp[0]
    L[temp2[1]-1] -= 1
    L[temp[1]-1] += 1
    cnt += 1
    ans = max(ans, num+cnt**2)

print(ans)

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)
que = []
res = 0
for i in range(K):
    res+=sushi[i][0]
    for j in sushi[i][1:]:
        que.append(j)
que.sort()
cur = res
res+=K**2
for var in range(1, K)[::-1]:
    a = que.pop()
    b = sushi[var][0]
    cur+=a-b
    res = max(res, cur+var**2)
print(res)

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

s = [tuple(map(int, input().split())) for _ in range(N)]
s.sort(key = lambda x: -x[1])
table = [(-1, -1)] * (N + 1)
c = 0
for i in range(N):
  if table[s[i][0]] == (-1, -1):
    table[s[i][0]] = (i, s[i][1])
    c += 1

table.sort(key = lambda x: -x[1])
res = [0] * (min(c, M) + 1)
vis = set()
p = M
for i in range(M): res[0] += s[i][1]
for x in range(1, min(c, M) + 1):
  res[x] = res[x - 1] + pow(x, 2) - pow(x - 1, 2)
  if table[x - 1][0] >= M:
    res[x] += table[x - 1][1]
    p -= 1
    while p in vis:
      p -= 1
    res[x] -= s[p][1]
  vis.add(table[x - 1][0])
  
print(max(res))

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

def main():
    N, K = map(int, input().split())
    d = defaultdict(list)
    for _ in range(N):
        t, v = map(int, input().split())
        d[t].append(v)

    other = []
    vmax = []
    for nums in d.values():
        nums.sort(reverse=True)
        for i, j in enumerate(nums):
            if i == 0:
                vmax.append(j)
            else:
                other.append(j)

    vmax.sort(reverse=True)
    other.sort(reverse=True)
    sum_v = [0] + list(accumulate(vmax))
    sum_o = [0] + list(accumulate(other))
    ans = 0
    for i in range(K,0,-1):
        neta = min(i, len(vmax))
        base = neta * neta
        tmp = base + sum_v[neta]
        if K - neta > len(other):
            break
        tmp += sum_o[K-neta]
        ans = max(ans, tmp)
        


    print(ans)

    
    




if __name__ == "__main__":
    main()

Solution (Python):


import sys

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10 ** 9)
INF = 10 ** 18
MOD = 10 ** 9 + 7

N, K = MAP()
DT = []
tmp = [[] for i in range(N)]
for i in range(N):
    t, d = MAP()
    DT.append((d, t))
DT.sort(reverse=1)

A = []
D = {}
dcnt = 0

for i in range(K):
    d, t = DT[i]
    dcnt += d
    
    if t not in D:
        D[t] = d
    
    else:
        if D[t] < d:
            D[t], d = d, D[t]
        A.append(d)
A.sort(reverse=1)
tcnt = len(D)

ans = dcnt + tcnt**2
for i in range(K, N):
    d2, t2 = DT[i]
    
    if t2 in D:
        continue
    D[t2] = d2
    
    if not A:
        break
    d = A.pop()
    
    dcnt -= d
    dcnt += d2
    tcnt += 1
    ans = max(ans, dcnt + tcnt**2)
print(ans)