qyliang/CodeNanoFix huggingface dataset:
Buggy code:
import sys
from heapq import *
input = sys.stdin.readline
N, K = map(int, input().split())
D, H, var, res, L = [], [], set(), 0, N.bit_length()+2
for _ in range(N):
    t,d = map(int, input().split())
    D.append((d<<L)+t)
D.sort(reverse=True)
for z in D[:K]:
    d, t = z>>L, z%L
    if t not in var: var.add(t)
    else: heappush(H, d)
    res += d
x = len(var)
res += x**2
tmp = res
for z in D[K:]:
    d, t = z>>L, z%L
    if H and t not in var:
        tmp += d-heappop(H)+2*x+1
        x += 1
        var.add(t)
        res = max(res,tmp)
print(res)

Fixed code:
import sys
from heapq import *
input = sys.stdin.readline
N, K = map(int, input().split())
D, H, var, res, L = [], [], set(), 0, N.bit_length()+2
for _ in range(N):
    t,d = map(int, input().split())
    D.append((d<<L)+t)
D.sort(reverse=True)
for z in D[:K]:
    d, t = z>>L, z%(1<<L)
    if t not in var: var.add(t)
    else: heappush(H, d)
    res += d
x = len(var)
res += x**2
tmp = res
for z in D[K:]:
    d, t = z>>L, z%(1<<L)
    if H and t not in var:
        tmp += d-heappop(H)+2*x+1
        x += 1
        var.add(t)
        res = max(res,tmp)
print(res)
Buggy code:
import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()


Fixed code:
import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            s.add(t)
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()
Buggy code:
from collections import defaultdict

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

data.sort(lambda x:-x[1])
KISO = 0
TOP = defaultdict(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if TOP[a] == 1:
    l -= 1
    continue
  if TOP[c] != 0:
    r += 1
    continue
  TOP[a] -= 1
  TOP[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Fixed code:
import sys
input = sys.stdin.readline
from collections import defaultdict

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

data.sort(key=lambda x:-x[1])
KISO = 0
EAT = defaultdict(int)
for i in range(K):
  t,v = data[i]
  EAT[t] += 1
  KISO += v
SHURUI = len(EAT)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if EAT[a] == 1:
    l -= 1
    continue
  if EAT[c] != 0:
    r += 1
    continue
  EAT[a] -= 1
  EAT[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  l -= 1
  r += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)
Buggy code:
N,K = map(int,input().split())
dic = {}
A = []
for _ in range(N):
    t,d = map(int,input().split())
    A.append([t,d])
A.sort(key = lambda x:x[1], reverse = True)
val = 0
set_ = set()
duplicated = collections.deque([])
for t,d in A[:K]:
    val += d
    if t in set_:
        duplicated.append(d)
    else:
        set_.add(t)
now = val + len(set_)**2

nx = collections.deque(A[K:])
ans = now
c = len(set_)
while len(set_) < K and nx:
    i,j = nx.popleft()
    if i in set_:
        continue
    d = duplicated.pop()
    diff = j-d + 2*c+1
    set_.add(i)
    now += diff
    ans =max(ans,now)
    c += 1
print(ans)


Fixed code:
import collections
N,K = map(int,input().split())
dic = {}
A = []
for _ in range(N):
    t,d = map(int,input().split())
    A.append([t,d])
A.sort(key = lambda x:x[1], reverse = True)
val = 0
set_ = set()
duplicated = collections.deque([])
for t,d in A[:K]:
    val += d
    if t in set_:
        duplicated.append(d)
    else:
        set_.add(t)
now = val + len(set_)**2

nx = collections.deque(A[K:])
ans = now
c = len(set_)
while len(set_) < K and nx:
    i,j = nx.popleft()
    if i in set_:
        continue
    d = duplicated.pop()
    diff = j-d + 2*c+1
    set_.add(i)
    now += diff
    ans =max(ans,now)
    c += 1
print(ans)
Buggy code:
import sys
input = sys.stdin.readline

N, K = map(int, input().split())
nums = []
for i in range(N):
    t, d = map(int, input().split())
    nums.append((d, t))
nums.sort(reverse=True)
se = set()
kind = 0
a1 = []
a2 = []
rest = []
for i, n in nums:
    if i < K:
        if not n[1] in se:
            se.add(n[1])
            kind += 1
            a1.append(n)
        else:
            a2.append(n)
    else:
        rest.append(n)
points = 0
for (d, t) in a1: points += d
for (d, t) in a2: points += d
ans = points + pow(kind, 2)
a2.sort()
for i, (d, t) in enumerate(rest):
    if t in se: continue
    if i >= len(a2): break
    points -= a2[i]
    points += d
    kind += 1
    ans = max(ans, points+pow(kind, 2))
print(ans)


Fixed code:
import sys
input = sys.stdin.readline

N, K = map(int, input().split())
nums = []
for i in range(N):
    t, d = map(int, input().split())
    nums.append((d, t))
nums.sort(reverse=True)
se = set()
a1 = []
a2 = []
rest = []
for i, n in enumerate(nums):
    if i < K:
        if not n[1] in se:
            se.add(n[1])
            a1.append(n)
        else:
            a2.append(n)
    else:
        rest.append(n)
points = 0
for (d, t) in a1: points += d
for (d, t) in a2: points += d
ans = points + pow(len(se), 2)
a2.sort()
idx2 = 0
for i, (d, t) in enumerate(rest):
    if t in se: continue
    if len(a2) <= idx2: break
    points -= a2[idx2][0]
    idx2 += 1
    points += d
    se.add(t)
    ans = max(ans, points+pow(len(se), 2))
print(ans)
Buggy code:
import sys
input = sys.stdin.buffer.readline

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

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

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

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

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

print(max(F))

Fixed code:
import sys
input = sys.stdin.buffer.readline

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

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

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

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

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

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

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

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

point=f+len(V)**2

ans=point

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

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

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

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

ans=point

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

print(ans)
Buggy code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]

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

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

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

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

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

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

print(ans)


Fixed code:
from heapq import heappush, heappop

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

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

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

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

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

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

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

print(ans)
Buggy code:
from heapq import heappop, heappush
N, K = map(int, input().split())
KMAX = 10 ** 5 + 1
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((d, t))

TD.sort(reverse=True)

pioneer = []
successor = []
appear = [False] * KMAX
cnt = 0
nsp = 0  # num of species
base = 0
best = []

for td in TD:
    d, t = td
    if cnt < K:
        heappush(successor if appear[t] else pioneer, (d, t))
        nsp += 0 if appear[t] else 1
        appear[t] = True
        base += d
        cnt += 1
        if cnt == K-1:
            heappush(best, -(base + nsp ** 2))
        continue
    if not successor:
        break
    if appear[t]:
        continue
    cd, ct = heappop(successor)
    appear[t] = True
    heappush(pioneer, (d, t))
    base += (d-cd)
    nsp += 1
    heappush(best, -(base + nsp ** 2))

print(-heappop(best))


Fixed code:
from heapq import heappop, heappush
N, K = map(int, input().split())
KMAX = 10 ** 5 + 1
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((d, t))

TD.sort(reverse=True)

pioneer = []
successor = []
appear = [False] * KMAX
cnt = 0
nsp = 0  # num of species
base = 0
best = []

for td in TD:
    d, t = td
    if cnt < K:
        heappush(successor if appear[t] else pioneer, (d, t))
        nsp += 0 if appear[t] else 1
        appear[t] = True
        base += d
        cnt += 1
        if cnt == K:
            heappush(best, -(base + nsp ** 2))
        continue
    if not successor:
        break
    if appear[t]:
        continue
    cd, ct = heappop(successor)
    appear[t] = True
    heappush(pioneer, (d, t))
    base += (d-cd)
    nsp += 1
    heappush(best, -(base + nsp ** 2))

print(-heappop(best))
Buggy code:
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)

Fixed code:
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()
cnt = 0
for t, d in td[:K]:
    cnt += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
cnt += len(v) ** 2
ans = cnt
for t, d in td[K:]:
    if t not in v and len(q) != 0:
        x = heappop(q)
        cnt += d + 2 * len(v) + 1 - x
        v.add(t)
        ans = max(ans, cnt)
print(ans)
Buggy code:
from heapq import heappop, heappush, heapify
n, k = map(int, input().split())
lst = [list(map(int, input().split())) for i in range(n)]
lst = sorted(lst, reverse=True, key=lambda x:x[1])
q = []
v = set()
cnt = 0
for t, d in lst[:k]:
    cnt += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
cnt += len(v)**2
ans = cnt
for t, d in lst[k:]:
    if t not in v and len(q)!=0:
        x = heappop(q)
        cnt += d + 2*len(v) + 1 - x
        ans = max(ans, cnt)
print(ans)

Fixed code:
from heapq import heappop, heappush
n, k = map(int, input().split())
lst = [list(map(int, input().split())) for i in range(n)]
lst = sorted(lst, reverse=True, key=lambda x:x[1])
q = []
v = set()
cnt = 0
for t, d in lst[:k]:
    cnt += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
cnt += len(v)**2
ans = cnt
for t, d in lst[k:]:
    if t not in v and len(q)!=0:
        x = heappop(q)
        cnt += d + 2*len(v) + 1 - x
        v.add(t)
        ans = max(ans, cnt)
print(ans)
Buggy code:
n, k = map(int, input().split())
TD = [[] for i in range(n)]
for i in range(n):
  TD[i] = list(map(int, input().split()))
  
TD.sort(key = lambda x:x[1], reverse = True)

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2

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

for i,j in TD[k:]:
  if i in su:
    continue
  else:
    su.add(i)
    temp = ans + j - huku.pop()[1] + len(su)**2 -(len(su)-1)**2
    ans = max(ans, temp)

print(ans)

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

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2
huku.sort(key = lambda x:x[1], reverse = True)

ansl = []
ansl.append(ans)
for i,j in TD[k:]:
  if i in su:
    continue
  elif huku == []:
    break
  else:
    temp = ansl[-1] + j - huku.pop()[1] + (len(su)+1)**2 -len(su)**2
    su.add(i)
    ansl.append(temp)

print(max(ansl))
Buggy code:
from heapq import heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()

for i in range(len(TD)):
    flg = 0
    if TD[i][0] not in muptiple:
        muptiple.add(TD[i][0])
        flg = 1
    TD[i][0] = flg
TD.reverse()

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = TD.pop()
    vari += v[0]
    point += v[1]
    if not v[0]:
        heappush(ret, t)

total = vari ** 2 + point

while len(TD) > 0 and len(ret) > 0:
    v = TD.pop()
    if v[0] == 0:
        continue
    pre_t = heappop(ret)
    point -= pre_t - v[1]
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)

Fixed code:
from heapq import heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()
hp = []
for t, d in TD:
    flg = 0
    if t not in muptiple:
        muptiple.add(t)
        flg = 1
    heappush(hp, (d*-1, flg))

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = heappop(hp)
    t = v[0]*-1
    vari += v[1]
    point += t
    if not v[1]:
        heappush(ret, t)

total = vari ** 2 + point

while len(hp) > 0 and len(ret) > 0:
    v = heappop(hp)
    if v[1] == 0:
        continue
    pre_t = heappop(ret)
    t = v[0]*-1
    point -= pre_t - t
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)
Buggy code:
from collections import defaultdict
from operator import itemgetter
import sys

input = sys.stdin.readline

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

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, key=itemgetter(0), reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)


Fixed code:
from collections import defaultdict
import sys

input = sys.stdin.readline

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

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)
Buggy code:
from collections import defaultdict

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

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

sushi.sort(key=lambda x: -x[1])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop(0)
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x

    if i == k - 1:
        x = len(kinds)
        ans_candidates[x] = umami + x * x
        remove_candidates.sort()
        print(umami)
        print(kinds)
        print(remove_candidates)

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))


Fixed code:
from collections import defaultdict

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

sushi.sort(key=lambda x: -x[1])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    # print("t, d = ", t, d)
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop()
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x
        # print("remove", remove_candidates)
        # print(x, umami, ans_candidates[x])

    if i == k - 1:
        x = len(kinds.keys())
        ans_candidates[x] = umami + x * x
        remove_candidates.sort(key=lambda x: -x)
        # print(umami)
        # print(kinds, x)
        # print(remove_candidates)
        # print(x, umami, ans_candidates[x])

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))
Buggy code:
from collections import defaultdict

def getlist():
	return list(map(int, input().split()))
D = defaultdict(int)
V = []

N, K = getlist()
for i in range(N):
	t, d = getlist()
	V.append([d, t])

V = sorted(V)

ans = []
ex = []
k = 0
for i in range(N):
	x = V[-1 - i][1]
	if k < K:
		if D[x] == 0:
			ans.append(V[-1 - i])
			D[x] += 1
			k += 1
	else:
		ex.append(V[-1 - i])


if k < K:
	for i in range(K - k):
		ans.append(ex[i])
	eex = []
	for i in range(K - k, len(ex)):
		eex.append(ex[i])
else:
	eex = ex
ans = sorted(ans)
eex = sorted(eex)


A = 0
for i in range(K):
	A += ans[i][0]
A += k * k
n = min(k, N - K)
for i in range(n):
	if ans[i][0] + ((k - i) ** 2) - ((k - i - 1) ** 2) < eex[-1 - i][0]:
		A += eex[-1 - i][0] - ans[i][0] - ((k - i) ** 2) + ((k - i - 1) ** 2)
	else:
		break

print(A)

Fixed code:
from collections import defaultdict

def getlist():
	return list(map(int, input().split()))
D = defaultdict(int)
V = []

N, K = getlist()
for i in range(N):
	t, d = getlist()
	V.append([d, t])

V = sorted(V)

ans = []
ex = []
k = 0
for i in range(N):
	x = V[-1 - i][1]
	if k < K:
		if D[x] == 0:
			ans.append(V[-1 - i])
			D[x] += 1
			k += 1
		else:
			ex.append(V[-1 - i])
	else:
		ex.append(V[-1 - i])


if k < K:
	for i in range(K - k):
		ans.append(ex[i])
	eex = []
	for i in range(K - k, len(ex)):
		eex.append(ex[i])
else:
	eex = ex
ans = sorted(ans)
eex = sorted(eex)


A = 0
for i in range(K):
	A += ans[i][0]
A += k * k
n = min(k, N - K)
for i in range(n):
	if ans[i][0] + ((k - i) ** 2) - ((k - i - 1) ** 2) < eex[-1 - i][0]:
		A += eex[-1 - i][0] - ans[i][0] - ((k - i) ** 2) + ((k - i - 1) ** 2)
	else:
		break

print(A)
Buggy code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r = []
    variety = set([])
    ans = 0
    for _c in choice:
        ans += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    ans += len(variety)**2
    i = k
    while len(variety) < min(max_variety, k):
        while td[i][0] in variety:
            i += 1
        variety.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))


Fixed code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r = []
    variety = set([])
    base = 0
    for _c in choice:
        base += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    base += len(variety)**2
    i = k
    ans = [base]
    while len(variety) < min(max_variety, k):
        while td[i][0] in variety:
            i += 1
        variety.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))