Problem:
 Score : 400 points    Problem Statement There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i .
You are choosing K among these N pieces to eat.
Your "satisfaction" here will be calculated as follows:   The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".  The base total deliciousness is the sum of the deliciousness of the pieces you eat.  The variety bonus is x*x , where x is the number of different kinds of toppings of the pieces you eat.   You want to have as much satisfaction as possible.
Find this maximum satisfaction.      Constraints  1 \leq K \leq N  \leq 10^5  1 \leq t_i \leq N  1 \leq d_i \leq 10^9  All values in input are integers.         Input Input is given from Standard Input in the following format:  N  K  t_1  d_1  t_2  d_2  .  .  .  t_N  d_N       Output Print the maximum satisfaction that you can obtain.        Sample Input 1 5 3
1 9
1 7
2 6
2 5
3 1      Sample Output 1 26  If you eat Sushi 1,2 and 3 :   The base total deliciousness is 9+7+6=22 .  The variety bonus is 2*2=4 .   Thus, your satisfaction will be 26 , which is optimal.       Sample Input 2 7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5      Sample Output 2 25  It is optimal to eat Sushi 1,2,3 and 4 .       Sample Input 3 6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000      Sample Output 3 4900000016  Note that the output may not fit into a 32 -bit integer type.

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

netano_min=min(S[j][0] for j in range(n))
for i in range(n):
    S[i][0]-=netano_min-1

neta=max(S[j][0] for j in range(n)) 

F=[0 for i in range(neta+1)] 

ma=sum(S[i][1] for i in range(k))

neta_s=len(set([S[i][0] for i in range(k)]))

netano_min=min(S[j][0] for j in range(n))

F[neta_s]=ma





netalist=[0 for i in range(n+1)] 
keseru=[] 
for i in range(k):
    if netalist[S[i][0]]>=1 :
        keseru.append(S[i])
        
    netalist[S[i][0]]+=1




right=len(keseru)-1
left=k
i=neta_s+1
while len(keseru)>0 and i<=k and right>=0 and left<=n-1 :
    while left<=n-1 and netalist[S[left][0]]!=0 :
        left+=1
    if left==n :
        break
    netalist[S[left][0]]+=1
    F[i]=F[i-1]+S[left][1]-keseru[right][1]
    right-=1
    left+=1
    i+=1


for j in range(neta_s,i+1-1) :
    F[j]+=j**2

print(max(F))

Solution (Python):
n, k = map(int, input().split())
netas_oishisas = []
for _ in range(n):
    netas_oishisas.append(tuple(map(int, input().split())))
netas_oishisas.sort(key = lambda no: no[1], reverse = True)
tabetaneta = set()
second = []
greedyscore = 0
for i in range(k):
    neta, oishisa = netas_oishisas[i]
    if neta in tabetaneta:
        second.append(oishisa)
    else:
        tabetaneta.add(neta)
    greedyscore += oishisa
greedyscore += len(tabetaneta) ** 2
maxscore = greedyscore
currentscore = greedyscore
second.reverse()
secondpointer = 0
secondn = len(second)
for i in range(k, n):
    if secondpointer >= secondn:
        break
    neta, oishisa = netas_oishisas[i]
    if neta not in tabetaneta:
        currentscore = currentscore - second[secondpointer] + 2*len(tabetaneta) + 1 + oishisa
        maxscore = max(maxscore, currentscore)
        secondpointer += 1
        tabetaneta.add(neta)
print(maxscore)

Solution (Python):
import sys

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    S = TD[:K]
    result = base_point + type_bonus**2
    
    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        if type_count_dict[s[0]] >= 2:
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            type_bonus += 1
            ridx += 1
            result = max(type_bonus**2 + base_point, result)
    print(result)


if __name__ == '__main__':
    main()

Solution (Python):

import sys



sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    
    
    
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    

    
    
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    

    S = TD[:K]
    
    
    
    

    
    
    
    result = base_point + type_bonus**2
    
    

    
    
    
    
    
    
    
    
    
    
    
    

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        
        
        if type_count_dict[s[0]] >= 2:
            
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            
            type_bonus += 1
            ridx += 1
            
            result = max(type_bonus**2 + base_point, result)

    
    print(result)


if __name__ == '__main__':
    main()

Solution (Python):
ri = lambda: int(input())
rl = lambda: list(map(int,input().split()))

N,K=rl()

td_1 = [0]*N
td_rem = []
for _ in range(N):
    t,d = rl()
    t -= 1
    if td_1[t] == 0:
        td_1[t] = d
    else:
        if d <= td_1[t]:
            td_rem += [d]
        else:
            td_rem += [td_1[t]]
            td_1[t] = d


td_1.sort(reverse=True)
td_rem.sort(reverse=True)


import numpy as np
index_0 = td_1.index(0) if 0 in td_1 else len(td_1)
td_1 = list(np.cumsum([0] + td_1[:index_0]))
td_rem = list(np.cumsum([0]+td_rem))

ans = [0 for _ in range(K+1)]
sum_l = 0
for i in range(1,K+1):
    if len(td_1)-1 >= i and len(td_rem)-1>=K-i:
        ans[i] = i*i + td_1[i] + td_rem[K-i]




print(max(ans))

Solution (Python):
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 top[t]>=d:
            sub.append(d)
        else:
            sub.append(top[t])
            top[t]=d
 
top.sort()
top = top[::-1]
res = [0]
for i in range(n):
    if top[i]!=0:
        res.append(res[-1]+top[i])
    else:
        res.append(-1)
sub.sort()
sub = sub[::-1]
w = [0]
for i in range(len(sub)):
    w.append(w[-1]+sub[i])
tmp = 0
for i in range(1,k+1):
    if len(w)-1>=k-i and res[i]!=-1:
        tmp = max(res[i]+w[k-i]+i*i,tmp)
 
print(tmp)

Solution (Python):
ri = lambda: int(input())
rl = lambda: list(map(int,input().split()))

N,K=rl()
td = [rl() for _ in range(N)]

from operator import itemgetter
key0 = itemgetter(0)
key1 = itemgetter(1)
td.sort(key=key1,reverse=True)

netas = set()
dup = []

for t,d in td[:K]:
    if t in netas:
        dup.append(d)
    else:
        netas.add(t)

score = len(netas)**2 + sum(map(key1,td[:K]))
ans = score
for t,d in td[K:]:
    if t not in netas and dup != []:
        netas.add(t)
        dup_ = dup.pop(-1)
        score += 2*len(netas)-1 + d - dup_
        ans = max(ans,score)
print(ans)

Solution (Python):
n, k = [int(i) for i in input().split(" ")]
t = []
d = []

for i in range(0, n):
    x, y = [int(i) for i in input().split(" ")]
    t.append(x)
    d.append(y)

dt = zip(d, t)
dt = sorted(dt, key = lambda k:k[0], reverse=True)
t_max = [False for i in range(0, n + 1)]
d_max = []
d_non_max = []

for i in range(0, n):
    if t_max[dt[i][1]]:
        d_non_max.append(dt[i][0])
    else:
        t_max[dt[i][1]] = True
        d_max.append(dt[i][0])
        
p = [0 for i in range(0, len(d_max) + 1)]
sum_max = 0
sum_non_max = sum(d_non_max[0:min(len(d_non_max), k)])

for i in range(1, len(p)):
    sum_max += d_max[i - 1]
    if k - i >= 0 and k - i < len(d_non_max):
        sum_non_max -= d_non_max[k - i]
    if k - i >= 0 and k - i - 1 < len(d_non_max):
        p[i] = sum_max + sum_non_max + i * i
    else:
        p[i] = 0
        
print(max(p))

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

removable_list = []
S = set()
base_pt = 0
for t, d in TD[:K]:
    if t in S:
        removable_list.append([t, d])
    else:
        S.add(t)
    base_pt += d
kind = len(S)
removable_list.sort(key=lambda x: x[1])

addable_list = []
for t, d in TD[K:]:
    if t in S:
        continue
    addable_list.append([t, d])
    S.add(t)

pt = [base_pt + kind*kind]
l = min(len(removable_list), len(addable_list))

for rl, al in zip(removable_list[:l], addable_list[:l]):
    kind += 1
    base_pt += al[1] - rl[1]
    pt.append(base_pt + kind*kind)
print(max(pt))

Solution (Python):
import heapq
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append((-d,t))
data=sorted(data)
neta=set()
que1=[]
p,x=0,0
for i in range(k):
    if data[i][1] in neta:
        heapq.heappush(que1,(-data[i][0],data[i][1]))
        x-=data[i][0]
    else:
        x-=data[i][0]
        p+=1
        neta.add(data[i][1])
mx=x+p*p
for i in range(k,n):
    if data[i][1] in neta:
        pass
    else:
        if que1:
            a=heapq.heappop(que1)
            x=x-a[0]-data[i][0]
            p+=1
            neta.add(data[i][1])
            mx=max(mx,x+p*p)
        else:
            break
print(mx)

Solution (Python):
N,K = list(map(int,input().split()))
TDlist = []

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

TDlist= sorted(TDlist, key = lambda x: x[1],reverse=True)
kinds = set()
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.add(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.add(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    answer = max(score,answer)

print(answer)

Solution (Python):
N, K = map(int, input().split())
a = [[int(t) for t in input().split()] for d in range(N)]
a.sort(key=lambda x: x[1], reverse=True)
x=[]
y=[]
st = set()

for i in range(K):
    t, d = a[i]
    if(t in st):
        x.append(d)
    else:
        y.append(d)
        st.add(t)
        
X = sum(x)
Y = sum(y)
V = len(y)
res = X + Y + V ** 2

for i in range(K, N):
    if not x:
        break
    t, d = a[i]
    if t in st:
        continue
    st.add(t)
    X -= x.pop()
    Y += d
    V += 1
    res = max(res, X + Y + V ** 2)
    
print(res)

Solution (Python):
def takeSecond(elem):
    return elem[1]

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

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

sushi.sort(key=takeSecond,reverse=True)
'''
selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])
maxscore = score
front_pointer = k - 1
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    
    for p in range(front_pointer,-1,-1):
        item = selected[p]
        flag = False
        for idx in range(p):
            if(selected[idx][0]==item[0]):
                score -= item[1]
                selected.remove(item)
                front_pointer = p - 1
                flag = True
                break
        if(flag):
            break

    
    for j in range(pointer,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            pointer = j + 1            
            score += 2*startKindNum + 1 + sushi[j][1]
            if(score>maxscore):
                maxscore = score
            startKindNum += 1
            break
'''
kindset = set([])
score = 0
extras = []
for item in sushi[:k]:
    if item[0] in kindset:
        extras.append(item[1])
    else:
        kindset.add(item[0])
    score += item[1]
score += len(kindset)**2
maxscore = score
for item in sushi[k:]:
    if not  extras:
        break
    if item[0] in kindset:
        continue
    score += (item[1] - extras.pop()) + 2*len(kindset) + 1
    kindset.add(item[0])
    
    if score > maxscore:
        maxscore = score


print(maxscore)

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

score = 0

remains = []
kinds = set()
for t, d in sushi[:K]:
    if t in kinds:
        remains.append(d)
    kinds.add(t)
    score += d
score += len(kinds)**2

test = score
for t, d in sushi[K:]:
    if not remains:
        break
    if t in kinds:
        continue
    kinds.add(t)
    rm = remains.pop()
    test = test + (d - rm) + (len(kinds)*2 - 1)
    score = max(score, test)

print(score)

Solution (Python):
import heapq as hq
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    hq.heappush(stack,[d,t])
ans = kiso + len(neta)**2
for t,d in t_d[K:]:
    if neta[t] > 0 : continue
    while(stack):
        d0,t0 = hq.heappop(stack)
        if neta[t0] > 1:
            kiso += d - d0
            neta[t0] -= 1
            neta[t] += 1
            ans = max(ans, kiso+len(neta)**2)
            hq.heappush(stack,[d0,t0])
            hq.heappop(stack)
            break
print(ans)

Solution (Python):
import heapq as hq
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    hq.heappush(stack,[d,t])
ans = kiso + len(neta)**2
for t,d in t_d[K:]:
    if neta[t] > 0 : continue
    while(stack):
        d0,t0 = hq.heappop(stack)
        if neta[t0] > 1:
            kiso += d - d0
            neta[t0] -= 1
            neta[t] += 1
            ans = max(ans, kiso+len(neta)**2)
            break
print(ans)

Solution (Python):
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    stack.append((t,d))
ans = kiso + len(neta)**2
tmp = ans
for t,d in t_d[K:]:
    if t in neta : continue
    while(stack):
        t0,d0 = stack.pop()
        if neta[t0] > 1:
            tmp += d - d0 + 2*len(neta) + 1
            ans = max(ans,tmp)
            neta[t0] -= 1
            neta[t] += 1
            break
    else:
        break
print(ans)

Solution (Python):
def solve(n, k, ts, ds, dt):
    firsts = set()
    seconds = []

    dt.sort(reverse=True)

    count = 0
    t_sum = 0
    d_sum = 0
    score_max = 0
    for d, t in dt:

        if count == k:
            
            if not seconds:
                break

            
            if t in firsts:
                continue

            d_sum -= seconds[-1]
            seconds.pop()
            count -= 1

        if t not in firsts:
            firsts.add(t)
            t_sum += 1
        else:
            seconds.append(d)

        d_sum += d
        score = d_sum + (t_sum * t_sum)
        

        if score_max < score:
            score_max = score

        count += 1

    return score_max


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

ts = [0] * n
ds = [0] * n
td = [(0, 0)] * n
for i in range(n):
    t, d = map(int, input().split())
    ts[i] = t
    ds[i] = d
    td[i] = (d, t)

print(solve(n, k, ts, ds, td))

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

A.sort(reverse=True)
X = []
Y = []
st = set()

for i in range(k):
    d,t = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)

x = sum(X)
y = sum(Y)
v = len(Y)
ans = x + y + v**2
for i in range(k, n):
    if not X:
        break
    d, t = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v**2)

print(ans)

Solution (Python):

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


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


X = [] 
Y = [] 
st = set() 
for i in range(k):
    t, d = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)


x = sum(X)
y = sum(Y)
v = len(Y) 
ans = x + y + v ** 2
for i in range(k, n):
    if not X:
        break
    t, d = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v ** 2)

print(ans)

Solution (Python):
from collections import defaultdict
import heapq


def get_first_try(K, ts, ds):
    eats = []
    non_eats = []
    for i in range(len(ts)):
        t = ts[i]
        d = ds[i]
        heapq.heappush(eats, (d, t))
        if len(eats) > K:
            non_eats.append(heapq.heappop(eats))
    return [(t, d) for d,t in eats], [(t, d) for d,t in non_eats]

def count_kinds(eats):
    kinds = defaultdict(int)
    for e in eats:
        kinds[e[0]] += 1
    return kinds

def calc_satisfaction(eats, eat_kinds):
    s = 0
    for e in eats:
        s += e[1]
    return s + calc_bonus(eat_kinds)


def calc_bonus(kinds):
    x = 0
    for k in kinds:
        if kinds[k] > 0:
            x += 1
    return x**2


def solve(N, K, ts, ds):
    eats, non_eats = get_first_try(K, ts, ds)
    eat_kinds = count_kinds(eats)
    satisfaction = calc_satisfaction(eats, eat_kinds)
    sorted_eats = sorted(eats, key=lambda x: x[1])
    sorted_non_eats = sorted(non_eats, key=lambda x: x[1], reverse=True)
    n_kinds = len(eat_kinds)

    max_satisfaction = satisfaction

    i = 0
    j = 0
    while i < len(sorted_eats) and j < len(sorted_non_eats):
        t, d = sorted_eats[i]
        if eat_kinds[t] == 1:
            i += 1
            continue
        t2, d2 = sorted_non_eats[j]
        if eat_kinds[t2] > 0:
            j += 1
            continue
        prev_bonus = n_kinds ** 2
        eat_kinds[t] -= 1
        eat_kinds[t2] += 1
        n_kinds += 1
        current_bonus = n_kinds ** 2
        satisfaction = satisfaction - d + d2 - prev_bonus + current_bonus
        if max_satisfaction < satisfaction:
            max_satisfaction = satisfaction
        i += 1
        j += 1

    return max_satisfaction


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

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

print(solve(N, K, ts, ds))

Solution (Python):
n, k = map(int, input().split())
sushi = []
for i in range(n):
    tmp = list(map(int, input().split()))
    sushi.append(tmp)
eaten = set() 
chohuku = [] 
scores = [] 

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

tmp = 0 
for i in range(k): 
    if sushi[i][0] in eaten: 
        chohuku.append(sushi[i][1])
    else:
        eaten.add(sushi[i][0]) 
    tmp += sushi[i][1] 
chohuku.sort()
scores.append(tmp+len(eaten)**2) 

for i in range(k,n):
    if not sushi[i][0] in eaten and chohuku != []:
        tmp -= chohuku.pop(0)
        eaten.add(sushi[i][0])
        tmp += sushi[i][1]
        scores.append(tmp+len(eaten)**2)
        
print(max(scores))

Solution (Python):
from collections import defaultdict

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

sushi = sorted(sushi, key = lambda x: x[1], reverse = True)
types = set()
max_vals = [0]
other_vals = [0]

for typ, val in sushi:
    if typ not in types:
        max_vals.append(max_vals[-1] + val)
        types.add(typ)
    else:
        other_vals.append(other_vals[-1] + val)

max_score = 0
for k in range(1, K + 1):
    if K - k < len(other_vals) and k < len(max_vals):
        score = max_vals[k] + other_vals[K - k] + k ** 2
        max_score = max(max_score, score)

print(max_score)

Solution (Python):
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
typ = set()
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.add(item[0])
        item[0] = 1
        lstA.append(item[1])
ans = 0
neta = 0
for i in range(K):
    ans += sushi[i][1]
    if sushi[i][0] == 1:
        neta += 1
ans += neta ** 2
C = sum(lstA[:neta])
D = sum(lstB[:K - neta])
for i in range(neta + 1, min(K, len(lstA)) + 1):
    C += lstA[i - 1]
    D -= lstB[K - i]
    tmp = C + D + i ** 2
    ans = max(ans, tmp)
print(ans)

Solution (Python):


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

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()

for i in range(k):
    op += sushi[i][1]

    t = sushi[i][0]
    if t_taberu.isdisjoint({t}):    
        t_taberu.add(t)
    else:
        pop_cands.append(i)

    taberu_ind.add(i)
mp_max = op + len(t_taberu)**2


i += 1
j = i
for i in range(i, n):
    if len(pop_cands) == 0:
        break

    pop_ind = pop_cands.pop()
    taberu_ind.remove(pop_ind)

    for j in range(j, n):
        t = sushi[j][0]
        if t_taberu.isdisjoint({t}):
            t_taberu.add(t)
            break
    else:
        break
    
    
    op -= sushi[pop_ind][1]
    op += sushi[j][1]
    mp = op + len(t_taberu)**2
    mp_max = max(mp_max, mp)

    taberu_ind.add(i)    
print(mp_max)

Solution (Python):
N, K = map( int, input().split())
T = [(0,0)]*N
for i in range(N):
    t, d = map( int, input().split())
    T[i] = (d,t)
T.sort( key= None, reverse = True)
C = T[:K]
AC = T[K:]
kiso = 0
s = 0
S = [0]*(N+1)
for i in range(K):
    d, t = T[i]
    kiso += d
    if S[t] == 0:
        S[t] = 1
        s += 1
    else:
        S[t] += 1
ans = kiso + s**2
C.sort()
r = 0
 
for i in range(K):
    d, t = C[i]
    if S[t] >= 2:
        while r < N-K:
            if S[ AC[r][1]] == 0:
                S[t] -= 1
                S[ AC[r][1]] += 1
                s += 1
                kiso -= d
                kiso += AC[r][0]
                r += 1
                break
            r += 1
    ans = max( ans, kiso + s**2)
print( ans)

Solution (Python):


from sys import stdin


input = stdin.readline

N, K = list(map(int, input().split()))
List1= []
kind = set()
point1, point2 = 0, 0

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

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

for ti, di in INPUT[:K]:
    point1 += di
    if ti in kind:
        List1.append(di)
    kind.add(ti)

point2 = len(kind)

Score = point1 + point2**2
Test = Score

List1POP = List1.pop

for ti, di in INPUT[K:]:
    if 0 >= len(List1):
        break
    if ti in kind:
        continue
    Test = Test - List1POP() + di + 2 * point2 + 1
    point2 += 1
    kind.add(ti)
    Score = max(Score, Test)

print(Score)

Solution (Python):

from sys import stdin
from collections import Counter
import numpy as np

s_in = lambda: stdin.readline()[:-1]  
d_in = lambda: int(stdin.readline())  
ds_in = lambda: list(map(int, stdin.readline().split()))  


N, K = ds_in()
t_list = []
d_list = []
for _ in range(N):
    t, d = ds_in()
    t_list.append(t)
    d_list.append(d)
np_t = np.array(t_list)  
np_d = np.array(d_list)  

sort_idx = np.argsort(np_d)[::-1][:K]  
top_sushi = np_t[sort_idx]
top_score = np_d[sort_idx]
kind = len(set(top_sushi))
point = top_score.sum() + kind**2  
point_list = [point]
sushi_counter = Counter(top_sushi)

sushi_set = {k: [] for k in set(t_list)}
for i, j in zip(t_list, d_list):
    sushi_set[i].append(j)
max_sushi = {k: max(sushi_set[k]) for k in set(t_list)}

for s in set(top_sushi):
    max_sushi.pop(s)  
max_values = list(max_sushi.values())
max_values.sort(reverse=True)
if len(max_values) > 0:
    for sushi, score in zip(top_sushi[::-1], top_score[::-1]):
        if sushi_counter[sushi] > 1:
            sushi_counter[sushi] -= 1
            kind += 1
            point = point - score + max_values.pop(0) + 2*kind - 1
            point_list.append(point)
            if len(max_values) == 0:
                break

print(max(point_list))

Solution (Python):
n,k=map(int,input().split())

td=[]

for i in range(n):
    a=[]
    a=[int(j) for j in input().split()]
    td.append(a)

from operator import itemgetter
td.sort(key=itemgetter(1))  
td.reverse()
td.sort(key=itemgetter(0))

s=0
l=[]
for i in range(n):
    if s!=td[i][0]:
        s=td[i][0]
        l.append([td[i][1],1])
    else:
        l.append([td[i][1],0])

l.sort(key=itemgetter(0))
l.reverse()

op=0
sp=0
ans=[]
sl=[]
l1=[]

for i in range(n):
    if i<k:
        op+=l[i][0]
        sp+=l[i][1]
        if l[i][1]==0:
            sl.append(l[i][0])
    else:
        if l[i][1]==1:
            l1.append(l[i][0])


ans=[op+sp**2]
sl.sort()
l1.sort()
l1.reverse()

for i in range(min(len(sl),len(l1))):
    ans.append(ans[-1]-sl[i]+l1[i]+sp*2+1)
    sp+=1

print(max(ans))

Solution (Python):
from collections import deque

if __name__ == '__main__':
    n, k = map(int, input().split())
    sushis = [tuple(int(s) for s in input().split()) for _ in range(n)]
    sushis.sort(key=lambda t: t[1], reverse=True)

    used_set = set()
    replace_stack = deque()
    add_queue = deque()
    ans = 0
    for i in range(k):
        t, d = sushis[i]
        if t in used_set:
            replace_stack.appendleft(d)
        else:
            used_set.add(t)
        ans += d
    num = len(used_set)
    ans += num * num

    for i in range(k, n):
        t, d = sushis[i]
        if t not in used_set:
            add_queue.append(d)
            used_set.add(t)

    tmp = ans
    for x, r, a in zip(range(num, k + 1), replace_stack, add_queue):
        tmp = tmp + 2 * x + 1 - r + a
        ans = max(ans, tmp)

    print(ans)

Solution (Python):



import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1], reverse=True)
  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = set()
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types.add(c[0])
  types = set(types)

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for l in leave:
    if len(types) >= types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1
      types.add(l[0])

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])


if __name__=='__main__':
  main()

Solution (Python):
import sys

def point(sushi, n, k):
  sushi = sorted(sushi, key=lambda x:-x[1])
  
  ate_t = set()
  dpoint = 0
  dup = []
  for s in sushi[:k]:
    t = s[0]
    d = s[1]
    if t in ate_t:
      dup.append(d)
    ate_t.add(t)
    dpoint += d
  
  ans = len(ate_t)**2 + dpoint
  
  for s in sushi[k:]:
    t = s[0]
    d = s[1]
    if len(dup) <= 0:
      break
    if t in ate_t:
      continue
    ate_t.add(t)
    rm = dup.pop()
    dpoint = dpoint - rm + d
    res = len(ate_t)**2 + dpoint
    ans = max(ans, res)
  
  return ans

def main():
  input = sys.stdin.readline
  n, k = map(int, input().split())
  sushi = []
  for _ in range(n):
    t, d = map(int, input().split())
    sushi.append([t,d])
  print(point(sushi, n, k))


if __name__ == '__main__':
  main()

Solution (Python):
class Sushi(object):
    def __init__(self, sushi_type, delicious):
        self.t = sushi_type
        self.d = delicious


def main() -> None:
    n, k = map(int, input().split())
    menu = [[int(x) for x in input().split()] for _ in range(n)]
    selected = []
    unselected = []
    cnt = 0
    p_satisfied = 0
    selected_type = set()

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

    for i in range(n):
        sushi = Sushi(menu[i][0], menu[i][1])
        if sushi.t not in selected_type and cnt < k:
            selected_type.add(sushi.t)
            cnt += 1
            selected += [menu[i]]
        else:
            unselected += [menu[i]]

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

    while cnt < k and len(unselected) > 0:
        p_satisfied += unselected[-1][1]
        unselected.pop()
        cnt += 1

    while True:
        t_cnt = len(selected)
        if len(selected) == 0 or len(unselected) == 0:
            break
        if t_cnt ** 2 + selected[-1][1] > (t_cnt - 1) ** 2 + unselected[-1][1]:
            break
        p_satisfied += unselected[-1][1]
        unselected.pop()
        selected.pop()

    for v in selected:
        p_satisfied += v[1]

    print(p_satisfied + len(selected) ** 2)


if __name__ == "__main__":
    main()