qyliang/CodeNanoFix huggingface dataset:
Buggy code:
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])

ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans - q.pop() - len(s) ** 2 + TD[k1][1] + (len(s)+1) ** 2
        ans = max(ans, ans1)
        if ans == ans1:
            s.add(TD[k1][0])
    k1 += 1
    
print(ans)

Fixed code:
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])
ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans1 + TD[k1][1] - q.pop() - len(s) ** 2 + (len(s) + 1) ** 2
        ans = max(ans, ans1)
        s.add(TD[k1][0])
    k1 += 1
print(ans)
Buggy code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
point=0
stack=[]
kind=set()
for t,d in TD[:K]:
    point+=d
    if t in kind:stack.append(d)
    else:kind.add(t)
ans1=delicious+len(kind)**2
ans2=ans1
for t,d in TD[K:]:
    if not stack:break
    if not t in kind:
        ans2+=d-stack.pop()+2*len(kind)+1
        kind.add(t)
        ans2=max(ans1,ans2)
print(ans2)

Fixed code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
point=0
stack=[]
kind=set()
for t,d in TD[:K]:
    point+=d
    if t in kind:stack.append(d)
    else:kind.add(t)
point=point+len(kind)**2
ans=point
for t,d in TD[K:]:
    if not stack:break
    if not t in kind:
        point+=d-stack.pop()+2*len(kind)+1
        kind.add(t)
        ans=max(ans,point)
print(ans)
Buggy code:
N,K = map(int,input().split())
lis = []
val = 0
ans = 0
s = set()
dub = []

for _ in range(N):
    lis.append(tuple(map(int,input().split())))

lis = sorted(lis,key=lambda x: x[1],reverse=True)

for t,v in lis[:K]:

    if t in s:
        dub.append(v)
    s.add(t)
    val += v

ans = val + len(s) ** 2

for t,v in lis[K:]:
    if len(dub) == 0:
        break
    if t in s:
        continue

    s.add(t)
    r = dub.pop()
    val += v - r
    ans = max(ans,val+len(s) ** 2)

    print(ans)

Fixed code:
N,K = map(int,input().split())
lis = []
val = 0
ans = 0
s = set()
dub = []

for _ in range(N):
    lis.append(tuple(map(int,input().split())))

lis = sorted(lis,key=lambda x: x[1],reverse=True)

for t,v in lis[:K]:

    if t in s:
        dub.append(v)
    s.add(t)
    val += v

ans = val + len(s) ** 2

for t,v in lis[K:]:
    if len(dub) == 0:
        break
    if t in s:
        continue

    s.add(t)
    r = dub.pop()
    val += v - r
    ans = max(ans,val+len(s) ** 2)

print(ans)
Buggy code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())

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

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

#-----------------------------------------
best_list=[]
best_list_neta=[]
not_best_list=[]

for i in li:
    if i[0] in best_list_neta:
        not_best_list.append(i[1])
    else:
        best_list.append(i[1])
        best_list_neta.append(i[0])


s2=min(len(best_list_neta),k)
s1=find_num_neta(list(range(k)),li)


score=-1
oishisa=[0 for i in range(n)]
oishisa[s1]=sum([li[p][1] for p in range(k)])

for i in range(s1+1,s2+1):
    oishisa[i]=oishisa[i-1]-not_best_list[k-i]+best_list[i-1]
    
    new_score=i**2+oishisa[i]

    if score<new_score:
        score=new_score
        
    
print(score)

Fixed code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())



def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))
#-----------------------------------------

li.sort(reverse=True,key=lambda x:x[1])
s2=min(find_num_neta(list(range(n)),li),k)
s1=find_num_neta(list(range(k)),li)
oishisa=sum([li[i][1] for i in range(k)])
score=oishisa+s1**2

#-----------------------------------------
li.sort(key=lambda x:x[0])
best_list=[]
not_best_list=[]


m=-1
for i in li:
    if m==i[0]:
        not_best_list.append(i)
    else:
        best_list.append(i)
        m=i[0]
        
not_best_list.sort(reverse=True,key=lambda x:x[1])
best_list.sort(reverse=True,key=lambda x:x[1])



 
 
for i in range(s1+1,s2+1):
    oishisa=oishisa+best_list[i-1][1]-not_best_list[k-i][1]
    
    new_score=i**2+oishisa
    
    if score<new_score:
        score=new_score
        
    
print(score)
Buggy code:
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)

Fixed code:
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)
Buggy code:
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=[]
que2=[]
p,x=0,0
for i in range(k):
    if data[i][1] in neta:
        heapq.heappush(que1,data[i])
        x-=data[i][0]
    else:
        x-=data[i][0]
        p+=1
        neta.add(data[i][1])
for i in range(k,n):
    if data[i][1]in neta:
        pass
    else:
        heapq.heappush(que2,data[i])
mx=x+p*p
while que1:
    a=heapq.heappop(que1)
    b==None
    while que2:
        b=heapq.heappop(que2)
        if b[1] not in neta:
            break
        b=None
    if b==None:
        break
    x=x+a[0]-b[0]
    p+=1
    neta.add(b[1])
    mx=max(mx,x+p*p)
print(mx)


Fixed code:
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)
Buggy code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sorted(sushi, key=lambda x: x[1])
sushi.reverse()
typ = []
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.append(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))):
    tmp = C + lstA[i] + D - lstB[K - i]
    tmp += i ** 2
    ans = max(ans, tmp)
print(ans)

Fixed code:
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)
Buggy code:
# -*- coding utf-8 -*-

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)    # sort by d

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()
# greedy for delicious point
for i in range(k):
    op += sushi[i][1]

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

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

# reduce dericious, gain variety
i += 1
j = i
for i in range(i, n):
    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
    
    # calc points
    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)



Fixed code:
# -*- coding utf-8 -*-

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)    # sort by d

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()
# greedy for delicious point
for i in range(k):
    op += sushi[i][1]

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

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

# reduce dericious, gain variety
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
    
    # calc points
    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)    # not necessary actually
print(mp_max)
Buggy code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  # sushis = sorted([tuple(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1])
  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 = []
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types += [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()

Fixed code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-

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()
Buggy code:
N, K = map(int, input().split())
menu = []

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

sorted(menu, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) + 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans) 

Fixed code:
N, K = map(int, input().split())
menu_ = []

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

menu = sorted(menu_, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta and duplicates:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) - 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans)
Buggy code:
n, k = map(int, input().split())
h = []
for i in range(n):
    t, d = map(int, input().split())
    h.append((t,d))
h.sort(key=lambda x: -x[1])

taste = 0
typen = 0
typecnt = {}
bad = []


for i in range(k):
    typecnt[h[i][0]] = typecnt.get(h[i][0], 0) + 1
    if(typecnt[h[i][0]] == 1):
        typen += 1
    else:
        bad.append(h[i][1])
    taste += h[i][1]


ans = taste + typen**2

for i in range(k, n):
    if(len(bad) == 0):
        break
    typen += 1
    taste -= bad.pop()
    taste += h[i][1]
    ans = max(ans, taste + typen**2)

print(ans)

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

taste = 0
typen = 0
typecnt = {}
bad = []


for i in range(k):
    typecnt[h[i][0]] = typecnt.get(h[i][0], 0) + 1
    if(typecnt[h[i][0]] == 1):
        typen += 1
    else:
        bad.append(h[i][1])
    taste += h[i][1]

# print(h)
# print(typecnt)
# print(bad)

ans = taste + typen**2

for i in range(k, n):
    if(len(bad) == 0):
        break
    cnt = typecnt.get(h[i][0], 0)
    if(cnt != 0):
        continue
    else:
        typecnt[h[i][0]] = cnt + 1
    typen += 1
    taste -= bad.pop()
    taste += h[i][1]
    ans = max(ans, taste + typen**2)

print(ans)
Buggy code:
N,K=map(int,input().split())

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0
number=[]
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.append(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[i][0] not in number:
    number.append(td[i][0])
    tempanswer=tempanswer-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)

Fixed code:
N,K=map(int,input().split())

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0
number=set()
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.add(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[index][0] not in number:
    number.add(td[index][0])
    tempanswer=tempanswer+td[index][1]-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)
Buggy code:
n, k = map(int, input().split())
dt = sorted(tuple(map(int, input().split()))[::-1] for _ in range(n))[::-1]
r = v = 0
a = []
s = set()
for i, (d,t) in enumerate(dt, 1):
    if t in s:
        if i <= k:
            a.append(d)
            v += d
    else:
        if i > k:
            v -= a.pop()
            if not a:
                break
        v += d
        s.add(t)
    r = max(r, len(s)**2+v)
print(r)

Fixed code:
n, k =map(int, input().split())
dt = sorted(tuple(map(int, input().split()))[::-1] for _ in range(n))[::-1]
r = v = 0
a = []
s = set()
for i, (d,t) in enumerate(dt, 1):
    if t in s:
        if i <= k:
            a.append(d)
            v += d
    else:
        if i > k:
            if not a:
                break
            v -= a.pop()
        v += d
        s.add(t)
    r = max(r, len(s)**2+v)
print(r)
Buggy code:
N,K=map(int,input().split())
td=[[int a for a in input().split()]for i in range(N)]

td=sorted(td,key=lambda x:-x[1])

ans=0
eat=[]
noteat=[]
for t,d in td[:K]:
  ans+=d
  if t not in eat:
    eat.append(t)
  else:
    noteat.append(d)
temp_result=ans+sum(eat)**2
max_result=temp_result

for t,d in td[K:]:
  if not noteat:
    break
  if t not in eat:
    temp_result+=2*len(eat)+1
    temp_result+=d-noteat.pop()
    eat.add(t)
    max_result=max(max_result,temp_result)
print(max_result)
 

Fixed code:
N,K=map(int,input().split())
td=[[int(a) for a in input().split()]for i in range(N)]

td=sorted(td,key=lambda x:-x[1])

ans=0
eat=set([])
noteat=[]
for t,d in td[:K]:
  ans+=d
  if t not in eat:
    eat.add(t)
  else:
    noteat.append(d)
temp_result=ans+len(eat)**2
max_result=temp_result

for t,d in td[K:]:
  if not noteat:
    break
  if t not in eat:
    temp_result+=2*len(eat)+1
    temp_result+=d-noteat.pop()
    eat.add(t)
    max_result=max(max_result,temp_result)
print(max_result)
Buggy code:
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()

print(A)

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

print(X)

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)


Fixed code:
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)
Buggy code:
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
n,k=IL()
top=dict()
sub=[]
for i in range(n):
    t,d=IL()
    if t in top:
        sub.append(min(top[t],d))
        top[t]=max(top[t],d)
    else:
        top[t]=d

top=sorted(top.items(),key=lambda x:-x[1])
sub.sort(reverse=True)

rtop=[top[0][1]]
rsub=[sub[0]]
for i in range(1,len(top)):
    rtop.append(rtop[-1]+top[i][1])
for i in range(1,len(sub)):
    rsub.append(rsub[-1]+sub[i])

ans=0
for i in range(k):
    if i<len(rtop) and 0<=k-i-2<len(rsub):
        ans=max(ans,rtop[i]+rsub[k-i-2]+(i+1)**2)
print(ans)

Fixed code:
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
n,k=IL()
top=dict()
sub=[]
for i in range(n):
    t,d=IL()
    if t in top:
        sub.append(min(top[t],d))
        top[t]=max(top[t],d)
    else:
        top[t]=d

top=sorted(top.items(),key=lambda x:-x[1])
sub.sort(reverse=True)

rtop=[top[0][1]]
rsub=[0]
for i in range(1,len(top)):
    rtop.append(rtop[-1]+top[i][1])
for i in range(len(sub)):
    rsub.append(rsub[-1]+sub[i])

ans=0
for i in range(k):
    if i<len(rtop) and 0<=k-i-1<len(rsub):
        ans=max(ans,rtop[i]+rsub[k-i-1]+(i+1)**2)
print(ans)
Buggy code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][0])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            ans = max(ans, num)


if __name__ == "__main__":
    main()


Fixed code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][1])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            ans = max(ans, num)
    print(ans)


if __name__ == "__main__":
    main()
Buggy code:
from collections import defaultdict
N, K = map(int, input().split())
T = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    T[t].append(d)
S = []
U = []
for k, v in T.items():
    v.sort(reverse=1)
    U.append((v[0], k))
    for w in v[1:]:
        S.append((w, k))
S.sort(reverse=1)
U.sort(reverse=1)
D = defaultdict(int)
cnt = 0
D[U[0][1]] += 1
t = 1
cnt += 1
res = U[0][0]
S = S[:K-1]
for w, k in S:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += w
ans = 0
for v, k in U[1:]:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += v
    while cnt > K:
        cnt -= 1
        w, l = S.pop()
        D[l] -= 1
        if D[l] == 0:
            t -= 1
        res -= w
    print(t, res, cnt)
    ans = max(ans, t*t + res)
print(ans)

Fixed code:
from collections import defaultdict
N, K = map(int, input().split())
T = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    T[t].append(d)
S = []; U = []
for k, v in T.items():
    v.sort(reverse=1)
    U.append((v[0], k))
    for w in v[1:]:
        S.append((w, k))
S.sort(reverse=1)
U.sort(reverse=1)
D = defaultdict(int)
t = cnt = 0
res = 0
for w, k in S:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += w
ans = 0
can = 1
for v, k in U:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += v
    while cnt > K:
        if not S:
            can = 0
            break
        cnt -= 1
        w, l = S.pop()
        D[l] -= 1
        if D[l] == 0:
            t -= 1
        res -= w
    if not can:
        break
    ans = max(ans, t*t + res)
print(ans)
Buggy code:
print('YNeos'[sorted(input().split())!=['1','4','7','9']::2])

Fixed code:
print('YNEOS'[sorted(input().split())!=['1','4','7','9']::2])