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
from heapq import heapify, heappop

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)

que = []
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]
        que.append((d, t))
heapify(que)
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
    tmp = []
    
    if not que:
        break
    d, t = heappop(que)
    
    dcnt -= d
    dcnt += d2
    tcnt += 1
    ans = max(ans, dcnt + tcnt**2)
print(ans)

Solution (Python):
n,k=map(int,input().split())
TD=[list(map(int,input().split())) for _ in range(n)]
TD=sorted(TD,key=lambda x:-x[1])

q=[]
ans=0
v=set()
cnt=0
for t,d in TD[:k]:
    ans+=d
    if t in v:
        q.append(d)
    else:
        v.add(t)

ans+=len(v)**2
cnt=ans
for t,d in TD[k:]:
    if t not in v and len(q)!=0:
        x=q.pop()
        ans+=d+2*len(v)+1-x
        v.add(t)
        cnt=max(ans,cnt)

print(cnt)

Solution (Python):
def solve(n, k, t, d):
    que = []
    used = set()
    base_score = 0
    for ti, di in sorted(zip(t, d), key=lambda _:-_[1])[:k]:
        if ti in used:
            que.append(di)
        else:
            used.add(ti)
        base_score += di
    num = len(used)
    best = base_score + num**2

    M = {}
    for ti, di in zip(t, d):
        if not ti in used:
            if not ti in M:
                M[ti] = 0
            M[ti] = max(M[ti], di)

    que = sorted(que)
    for i, p in enumerate(sorted(M.values())[::-1][:len(que)]):
        base_score += p - que[i]
        num += 1
        best = max(best, base_score + num**2)
    return best

n, k = map(int, input().split())
t = [0] * n
d = [0] * n
for i in range(n):
    t[i], d[i] = map(int, input().split())
print(solve(n, k, t, d))

Solution (Python):


import sys

N, K = map(int, sys.stdin.buffer.readline().split())
td = []
for _ in range(N):
    tmp1, tmp2 = map(int, sys.stdin.buffer.readline().split())
    td.append([tmp1-1, tmp2, -1])

td = sorted(td, key = lambda x: x[1], reverse = True)
t_counter = [0]*N
for idx in range(N):
    t_counter[td[idx][0]] += 1
    td[idx][2] = t_counter[td[idx][0]]

t_set = set([td[idx][0] for idx in range(K)])
ans_list = []
ans = sum([td[idx][1] for idx in range(K)]) + len(t_set)**2
ans_list.append(ans)

a = 0
b = 0
while 1:
    while K - 1 - a >= 0:
        if td[K - 1 - a][2] >= 2:
            ans -= td[K - 1 - a][1]
            a += 1
            break
        a += 1
    if K-1-a < 0:
        break
    while K + b <= N-1:
        if td[K + b][0] not in t_set:
            ans -= len(t_set)**2
            t_set.add(td[K + b][0])
            ans += td[K + b][1] + len(t_set)**2
            b += 1
            break
        b += 1
    ans_list.append(ans)
    if K+b == N:
        break

print(max(ans_list))

Solution (Python):
def solve():
    from collections import deque, namedtuple
    from operator import attrgetter
    import sys

    input = sys.stdin.readline
    Sushi = namedtuple('Sushi', 'taste kind')

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

    e = []
    for _ in range(n):
        t, d = map(int, input().split())
        e.append(Sushi(taste=d, kind=t - 1))

    e.sort(key=attrgetter('taste'), reverse=True)

    ret = 0
    used = [False] * n  
    q = deque()  
    x = 0  
    for s in e[:k]:  
        if used[s.kind]:
            q.append(s)  
        else:
            used[s.kind] = True
            x += 1
        ret += s.taste
    ret += pow(x, 2)  

    t = ret
    r = iter(e[k:])  
    while q:
        rem = q.pop()  
        t -= rem.taste
        for s in r:
            if used[s.kind]:
                continue
            used[s.kind] = True
            t += s.taste + x * 2 + 1
            x += 1
            ret = max(ret, t)
            break
        else:
            break

    return ret


if __name__ == '__main__':
    print(solve())

Solution (Python):
def solve():
    from collections import namedtuple
    from heapq import heappush, heappop
    from operator import attrgetter
    import sys

    input = sys.stdin.readline
    Sushi = namedtuple('Sushi', 'taste kind')

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

    e = []
    for _ in range(n):
        t, d = map(int, input().split())
        e.append(Sushi(taste=d, kind=t - 1))

    e.sort(key=attrgetter('taste'), reverse=True)

    b = e[:k]  

    ret = 0
    used = [False] * n  
    h = []  
    x = 0  
    for s in b:
        if used[s.kind]:
            heappush(h, s)
        else:
            used[s.kind] = True
            x += 1
        ret += s.taste

    ret += pow(x, 2)  

    t = ret
    r = iter(e[k:])  
    while h:
        rem = heappop(h)  
        t -= rem.taste

        for s in r:
            if used[s.kind]:
                continue
            used[s.kind] = True
            t += s.taste + x * 2 + 1
            x += 1
            ret = max(ret, t)
            break
        else:
            break

    return ret


if __name__ == '__main__':
    print(solve())

Solution (Python):



from heapq import heappush, heappop
import sys

input = sys.stdin.readline

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

for lis in kind_value:
    lis.sort()

kind_value = sorted(
    (lis for lis in kind_value if lis),
    key=lambda x: x[-1], reverse=True)

ret = 0
curr = 0
kind = 0
remains = []
for lis in kind_value:
    kind += 1
    if kind > k:
        break
    curr += lis[-1]
    for d in lis[:-1]:
        heappush(remains, d)
        curr += d
    while len(remains) > k - kind:
        curr -= heappop(remains)
    ret = max(ret, curr + kind * kind)

print(ret)

Solution (Python):



from heapq import heappush, heappop
import sys

input = sys.stdin.readline

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

for lis in kind_value:
    lis.sort()

kind_value = sorted(
    (lis for lis in kind_value if lis),
    key=lambda x: x[-1], reverse=True)

ret = 0
curr = 0
kind = 0
remains = []
for lis in kind_value:
    kind += 1
    if kind > k:
        break
    curr += lis[-1]
    for d in lis[:-1]:
        heappush(remains, d)
        curr += d
    while len(remains) > k - kind:
        curr -= heappop(remains)
    ret = max(ret, curr + kind * kind)

print(ret)

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


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

sushi.sort(key=lambda x: x[1])
kinds = defaultdict(int)
eaten = []
ans = 0
for i in range(k):
    neta, point = sushi.pop()
    kinds[neta] += 1
    ans += point
    heapq.heappush(eaten, [point, neta])

num_of_kinds = len(kinds.keys())
ans += num_of_kinds ** 2
candidate = ans

for i in range(k):
    point, neta = heapq.heappop(eaten)
    if kinds[neta] == 1:
        continue
    kinds[neta] -= 1

    while sushi:
        new_neta, new_point = sushi.pop()
        if kinds[new_neta] >= 1:
            continue
        else:
            break
    else:
        break
    kinds[new_neta] += 1
    num_of_kinds += 1
    delta = new_point - point + (2*num_of_kinds - 1)
    candidate += delta
    ans = max(ans, candidate)

print(ans)

Solution (Python):
from collections import defaultdict
import heapq


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

sushi.sort(key=lambda x: x[1])
kinds = defaultdict(int)
eaten = []
ans = 0
for i in range(k):
    neta, point = sushi.pop()
    kinds[neta] += 1
    ans += point
    heapq.heappush(eaten, [point, neta])

num_of_kinds = len(kinds.keys())
ans += num_of_kinds ** 2
candidate = ans

for i in range(k):
    point, neta = heapq.heappop(eaten)
    if kinds[neta] == 1:
        continue
    kinds[neta] -= 1

    while sushi:
        new_neta, new_point = sushi.pop()
        if kinds[new_neta] >= 1:
            continue
        else:
            break
    else:
        break
    kinds[new_neta] += 1
    num_of_kinds += 1
    delta = new_point - point + (2*num_of_kinds - 1)
    candidate += delta
    ans = max(ans, candidate)

print(ans)

Solution (Python):
import sys
from operator import itemgetter
read = sys.stdin.read

N, K, *td = map(int, read().split())
td = sorted(zip(*[iter(td)] * 2), key=itemgetter(1), reverse=True)

stack = []
taste = 0
neta_group = set()
for t, d in td[:K]:
    taste += d
    if t in neta_group:
        stack.append(d)
    else:
        neta_group.add(t)

variety = len(neta_group)
candidates = [taste + variety ** 2]

if N == K:
    print(candidates[0])
    exit()

rest = list(reversed(td[K::]))

while rest and stack:
    t1, d1 = rest.pop()
    if t1 not in neta_group:
        neta_group.add(t1)
        d2 = stack.pop()
        variety += 1
        taste = taste - d2 + d1

        candidates.append(taste + variety ** 2)

print(max(candidates))

Solution (Python):
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
selected0 = set()
selected1 = set()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        if t in selected0:
            heapq.heappush(ls0, (d, t))
        selected0.add(t)
        base += d
    else:
        if t not in selected0 and t not in selected1:
            heapq.heappush(ls1, (-d, t))
            selected1.add(t)

m = len(selected0)
res = base + m**2

for _ in range(k):
    if len(ls0) == 0 or len(ls1) == 0:
        break
    d, _ = heapq.heappop(ls0)
    base -= d
    d, _ = heapq.heappop(ls1)
    base += -d
    m += 1
    res = max(res, base + m**2)
print(res)

Solution (Python):
import collections
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
counter0 = collections.Counter()
counter1 = collections.Counter()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        if counter0[t] > 0:
            heapq.heappush(ls0, (d, t))
        counter0[t] += 1
        base += d
    else:
        if t not in counter0 and t not in counter1:
            heapq.heappush(ls1, (-d, t))
            counter1[t] += 1

m = len(counter0)
res = base + m**2

for _ in range(k):
    if len(ls0) == 0 or len(ls1) == 0:
        break
    d, _ = heapq.heappop(ls0)
    base -= d
    d, _ = heapq.heappop(ls1)
    base += -d
    m += 1
    res = max(res, base + m**2)
print(res)

Solution (Python):
import collections
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
counter = collections.Counter()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        heapq.heappush(ls0, (d, t))
        counter[t] += 1
        base += d
    else:
        if t not in counter:
            heapq.heappush(ls1, (-d, t))

m = len(counter)
res = base + m**2

for i in range(k):
    ok = False
    while ls0:
        d, t = heapq.heappop(ls0)
        counter[t] -= 1
        if counter[t] > 0:
            base -= d
            ok = True
            break
    if not ok:
        break

    ok = False
    while ls1:
        d, t = heapq.heappop(ls1)
        if t not in counter:
            counter[t] += 1
            base += -d
            ok = True
            break
    if not ok:
        break

    res = max(res, base + (m + i + 1)**2)
print(res)

Solution (Python):
from collections import defaultdict

N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda A: A[1], reverse=True)

A = TD[:K][:: -1]
cnt = defaultdict(int)
V = set()
dSum = 0
for t, d in A:
    dSum += d
    cnt[t] += 1
    V.add(t)

i = 0
ans = dSum + len(V)**2
for t, d in TD[K:]:
    if not t in V:
        while i < K and cnt[A[i][0]] <= 1:
            i += 1
        if i == K:
            break

        dSum = dSum - A[i][1] + d
        V.add(t)
        cnt[A[i][0]] -= 1
        i += 1
        ans = max(ans, dSum + len(V)**2)

print(ans)

Solution (Python):
n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x: -x[1])

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    type, taste = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        same.append(taste)

rem = []
typesRem = set()
for i in range(nTake, n):
    type, taste = a[i]
    if type not in types and type not in typesRem:
        typesRem.add(type)
        rem.append(taste)

same.sort()
rem.sort(key=lambda x: -x)
ret = totalTaste + len(types) ** 2
curTotal = totalTaste
for i in range(min(len(same), len(rem))):
    curTotal = curTotal - same[i] + rem[i]
    ret = max(ret, curTotal + (len(types) + i + 1) ** 2)
print(ret)

Solution (Python):
n, k = map(int, input().split())
td = list(sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1]))
td2 = []
eaten = set()
dup = 0
score = k * k
i = 0
j = 0
while len(eaten) + dup < k:
    if i < n and td[i][0] in eaten:
        td2.append(td[i])
        i += 1
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if i < n and (j >= len(td2) or td[i][1] >= td2[j][1] - diff):
            t, d = td[i]
            score += d
            eaten.add(t)
            i += 1
        else:
            t, d = td2[j]
            score += d - diff
            dup += 1
            j += 1
print(score)

Solution (Python):
from heapq import heappop, heappush

n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x: (-x[1], x[0]))

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    type, taste = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        heappush(same, taste)

rem = []
typesRem = set()
for i in range(nTake, n):
    type, taste = a[i]
    if type not in types and type not in typesRem:
        typesRem.add(type)
        heappush(rem, -taste)

ret = totalTaste + len(types) ** 2
curTotal = totalTaste
add = 1
while same and rem:
    worst = heappop(same)
    best = heappop(rem)
    curTotal = curTotal - worst - best
    ret = max(ret, curTotal + (len(types) + add) ** 2)
    add += 1
print(ret)

Solution (Python):
from collections import*
n,k=map(int,input().split())
sushi=sorted([tuple(map(int,input().split()))for i in range(n)],key=lambda x:x[1],reverse=1)
first=defaultdict(list)
second=defaultdict(list)
menu=set()

for _,v in sushi[:k]:
    menu.add(_)
    first[_].append(v)
second=sushi[k:][::-1]
score=sum(map(sum,first.values()))
change=[]
for _,v in first.items():
    for i in v[1:]:
        change.append((_,i))
change.sort(key=lambda x:x[1],reverse=1)


ans=[score+len(first)**2]
while change and second:
    _,s=change.pop()
    score-=s
    try:
        while second[-1][0] in menu:
            pp,ss=second.pop()
        pp,ss=second.pop()
        menu.add(pp)
        score+=ss
        ans.append(score+len(menu)**2)
    except:
        break
print(max(ans))

Solution (Python):
from heapq import heappush, heappop
N, K = map(int, input().split())
log = [0]*N
pre = []
ls = []
for i in range(N):
  t, d = map(int, input().split())
  t = t-1
  heappush(ls,(-d,t))
v = 0
n = 0
for i in range(K):
  d, t = heappop(ls)
  d = -d
  log[t] += 1
  v += d
  if log[t]>1:
    heappush(pre,(d,t))
  else:
    n += 1
v += n*n
ans = v
while ls and pre:
  d, t = heappop(ls)
  if log[t]>0:
    continue
  log[t] += 1
  d = -d
  pv = v
  x, y = heappop(pre)
  v -= x
  v += d
  v += 2*n+1
  n += 1
  if ans<v:
    ans = v
print(ans)

Solution (Python):
import sys
sys.setrecursionlimit(2147483647)
INF=float("inf")
MOD=10**9+7
input=lambda :sys.stdin.readline().rstrip()
def resolve():
    from heapq import heappop,heappush
    n,k=map(int,input().split())
    A=[0]*n
    for i in range(n):
        t,d=map(int,input().split())
        t-=1
        A[i]=(t,d)
    A.sort(lambda x:-x[1]) 

    
    B=[0]*n
    first=[1]*n
    for i,a in enumerate(A):
        t,d=a
        if(first[t]):
            first[t]=0
            B[i]=(d,1)
        else:
            B[i]=(d,0)

    
    S=[]
    for i in range(k):
        d,f=B[i]
        heappush(S,(f,d))
    x=sum(S[i][1] for i in range(k))
    y=sum(S[i][0] for i in range(k))
    ans=x+y*y

    
    B=B[k:]
    for i in range(n-k):
        d,f=B[i]
        if(f==0): continue
        pf,pd=heappop(S)
        if(pf==1): break
        x+=d-pd
        y+=1
        ans=max(ans,x+y*y)
    print(ans)
resolve()

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)
selected = TD[:K]
dlc = sum(map(itemgetter(1), selected))
variety = set(map(itemgetter(0), selected))

used = set()
overlap = [] 
for t, d in TD[:K]:
    if t not in used:
        used.add(t)
    else:
        overlap.append(d)
answer = dlc + len(variety)**2

for t, d in TD[K:]:
    if not overlap:
        break
    if t in used:
        continue
    ov = overlap.pop()
    dlc = dlc - ov + d
    used.add(t)
    answer = max(answer, dlc + len(used)**2)

print(answer)

Solution (Python):
from heapq import heappush, heappop

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

counter = [0] * (N + 1)
h = []
for t, d in td :
    counter[t] += 1
    heappush(h, (d, d, t))

kind = len(set(t for t, d in td))
for _ in range(N - K) :
    while True :
        _, d, t = heappop(h)
        if counter[t] == 1 :
            loss = kind * kind - (kind - 1) * (kind - 1) + d
            if loss <= h[0][0] :
                counter[t] -= 1
                kind -= 1
            else :
                heappush(h, (loss, d, t))
                continue
        else :
            counter[t] -= 1
        
        break
        
print(sum(d for _, d, t in h) + kind * kind)

Solution (Python):






from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []

for i in range(k):
    d,t = sushi[i]
    heapq.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 = heapq.heappop(sushiq)
    
    while sushidict[t_min] == 1 and sushiq:
        d_min, t_min = heapq.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 == n_kind_all:
        break

print(max(res))

Solution (Python):
n,k=map(int,input().split())
TD=[list(map(int,input().split())) for _ in range(n)]
TD=sorted(TD,key=lambda x:-x[1])

q=[]
ans=0
v=set()
cnt=0
for t,d in TD[:k]:
    ans+=d
    if t in v:
        q.append(d)
    else:
        v.add(t)

ans+=len(v)**2
cnt=ans
for t,d in TD[k:]:
    if t not in v and len(q)!=0:
        x=q.pop()
        ans+=d+2*len(v)+1-x
        v.add(t)
        cnt=max(ans,cnt)

print(cnt)

Solution (Python):
n,k=map(int,input().split())
a=[]
for _ in range(n):a.append(tuple(list(map(int,input().split()))[::-1]))
a.sort(reverse=1)
x,y=[],[]
s=set()
for i in range(k):
  d,t= a[i]
  if t in s:
    x.append(d)
  else:
    y.append(d)
    s.add(t)
xx,yy,v=sum(x),sum(y),len(y)
ans=xx+yy+v**2
for i in range(k,n):
    if len(x)==0:break
    d,t=a[i]
    if t in s:continue
    s.add(t)
    xx-=x.pop()
    yy+=d
    v+=1
    ans=max(ans,xx+yy+v**2)
print(ans)

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

def main():
    n, k = map(int, input().split())
    td = [tuple(map(int, input().split())) for _ in range(n)]
    td = list(sorted(td, key=lambda x: -x[1]))
    dsum, kind = 0, 0
    num = [0 for _ in range(n)]
    h = []
    for t, d in td[:k]:
        dsum += d
        if num[t-1] > 0:
            heappush(h, (d, t))
        if num[t-1] == 0:
            kind += 1
        num[t-1] += 1
    
    ans = dsum + kind * kind

    for t, d in td[k:]:
        if len(h) == 0:
            break
        if num[t-1] > 0:
            continue
        else:
            _d, _t = heappop(h)
            num[_t-1] -= 1
            num[t-1] += 1
            kind += 1
            dsum += d-_d
            if dsum + kind * kind > ans:
                ans = dsum + kind * kind
    
    print(ans)


if __name__ == '__main__':
    main()

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

l=[] 
m=[] 
s=set()

for i in range(N):
    t,d=map(int,input().split())
    l.append((d,t))
l.sort(reverse=True) 

ans=K**2
t=0 

i=0 
j=0 

for _ in range(K): 
    while i<N and l[i][1] in s:
        m.append(l[i]) 
        i+=1 
    if (i<N and (m and j<len(m)) and l[i][0]>=m[j][0]-(2*K-1-2*t)) or (not m or j>=len(m)):
        ans+=l[i][0]
        s.add(l[i][1])
        i+=1
    else:
        ans+=m[j][0]-(2*K-1-2*t)
        j+=1
        t+=1

print(ans)

Solution (Python):
from heapq import heappop, heappush
 
N, K = map(int, input().split())
td = sorted([list(map(int, input().split()))
             for _ in range(N)], key=lambda x: x[1], reverse=True)

q = []
v = set()
count = 0
for t, d in td[:K]:
    count += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
count += len(v) ** 2
ans = count
for t, d in td[K:]:
    if t not in v and len(q) != 0:
        x = heappop(q)
        count += d + 2 * len(v) + 1 - x
        v.add(t)
        ans = max(ans, count)
        
print(ans)

Solution (Python):
import sys
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append([d,t])
data.sort(reverse=True)
flag=[0]*(n+1)
eat=[]
lst=[]
count=0
for u in data:
    d,t=u
    if flag[t]==0:
        eat.append(u)
        flag[t]=1
        count+=1
        if count==k:
            break
    else:
        lst.append(u)
else:
    lst.sort(reverse=True)
    sum=0
    for i in range(n-count):
        if i<k-count:
            sum+=lst[i][0]
        else:
            if lst[i][0]>eat[count-1][0]+2*count-1:
                eat[count-1]=lst[i]
                count-=1
            else:
                break
    for u in eat:
        sum+=u[0]
    sum+=count**2
    print(sum)
    sys.exit()
lst.sort(reverse=True)
for u in lst:
    d=u[0]
    if d>eat[count-1][0]+2*count-1:
        eat[count-1]=u
        count-=1
    else:
        break
sum=0
for u in eat:
    sum+=u[0]
sum+=count**2
print(sum)

Solution (Python):
import heapq

INF=10**12

N, K = map(int,input().split())
td = [[0,0] for i in range(N)]
s = [[] for i in range(N)]

for i in range(N):
    td[i][0], td[i][1] = map(int,input().split())
    td[i][0] -= 1
    s[td[i][0]].append(td[i][1])

for i in range(N):
    if s[i] == []:
        s[i].append(-INF)
    s[i].sort(key=lambda x:-x)

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

tmp = 0
q = []

for i in range(K):
    tmp += s[i][0]
    for j in range(1,len(s[i])):
        heapq.heappush(q,-s[i][j])
for i in range(K,N):
    for j in range(len(s[i])):
        heapq.heappush(q,-s[i][j])

res = tmp + K * K

for i in range(1,K)[::-1]:
    v = s[i][0]
    if q == []: break
    w = - q[0]
    if v < w:
        tmp = tmp + w - v
        heapq.heappop(q)
        heapq.heappush(q,-v)
    res = max(res, tmp + i * i)        

print(res)

Solution (Python):
from heapq import heappush, heapreplace

def main():
    INF = 10 ** 9 + 5

    N, K, *TD = map(int, open(0).read().split())

    E = [[] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        if e:
            e.sort(reverse=True)
        else:
            e.append(-INF)
    E.sort(key=lambda x: -x[0])

    cur = 0
    Q = []
    for ei in E[:K]:
        cur += ei[0]
        for eij in ei[1:]:
            heappush(Q, -eij)

    for ei in E[K:]:
        for eij in ei:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        tmp = cur + x * x
        if res < tmp:
            res = tmp

    print(res)

if __name__ == '__main__':
    main()

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

def main():
    INF = 10 ** 9 + 7

    N, K, *TD = map(int, open(0).read().split())

    E = [[] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        if e:
            e.sort(reverse=True)
        else:
            e.append(-INF)
    E.sort(key=lambda x: -x[0])

    cur = 0
    Q = []
    for i in range(K):
        cur += E[i][0]
        for j in range(1, len(E[i])):
            heappush(Q, -E[i][j])

    for i in range(K, N):
        for eij in E[i]:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        tmp = cur + x * x
        if res < tmp:
            res = tmp

    print(res)

if __name__ == '__main__':
    main()

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

def main():
    INF = 10 ** 9 + 7

    N, K, *TD = map(int, open(0).read().split())

    E = [[] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        if e:
            e.sort(reverse=True)
        else:
            e.append(-INF)
    E.sort(reverse=True)

    cur = 0
    Q = []
    for i in range(K):
        cur += E[i][0]
        for j in range(1, len(E[i])):
            heappush(Q, -E[i][j])

    for i in range(K, N):
        for eij in E[i]:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        tmp = cur + x * x
        if res < tmp:
            res = tmp

    print(res)

if __name__ == '__main__':
    main()

Solution (Python):
from heapq import heappush, heapreplace

def main():
    INF = 10 ** 9 + 5

    N, K, *TD = map(int, open(0).read().split())

    E = [[-INF] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        e.sort(reverse=True)
    E.sort(key=lambda x: -x[0])

    cur = 0
    Q = []
    for ei in E[:K]:
        cur += ei[0]
        for eij in ei[1:]:
            heappush(Q, -eij)

    for ei in E[K:]:
        for eij in ei:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        res = max(res, cur + x * x)

    print(res)

if __name__ == '__main__':
    main()

Solution (Python):
from heapq import heappush, heapreplace

def main():
    INF = 10 ** 9 + 5

    N, K, *TD = map(int, open(0).read().split())

    E = [[] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        if e:
            e.sort(reverse=True)
        else:
            e.append(-INF)
    E.sort(key=lambda x: -x[0])

    cur = 0
    Q = []
    for ei in E[:K]:
        cur += ei[0]
        for eij in ei[1:]:
            heappush(Q, -eij)

    for ei in E[K:]:
        for eij in ei:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        res = max(res, cur + x * x)

    print(res)

if __name__ == '__main__':
    main()