Buggy python submission (verdict: Runtime Error):
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []
for key in sushi:
    sushi[key].sort(reverse=True)
    first.append(sushi[key][0])
    second += sushi[key][1:]

first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))


ans = 0
for i in range(1, K+1):
    if K - i <= len(second_cum)-1:
        ans = max(ans, first_cum[i] + second_cum[K-i] + i*i)

print(ans)

Fixed submission:
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []

for nums in sushi.values():
    nums.sort(reverse=True)
    for i, j in enumerate(nums):
        if i == 0:
            first.append(j)
        else:
            second.append(j)
first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))

ans = 0

for i in range(K, 0, -1):
    neta = min(i, len(first))
    tmp = neta * neta + first_cum[neta]

    if K - neta > len(second):
        break
    
    tmp += second_cum[K - neta]

    ans = max(ans, tmp)

print(ans)

Buggy python submission (verdict: Time Limit Exceeded):
import sys
input = lambda: sys.stdin.readline().rstrip()

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

sushi = [[] for _ in range(N)]

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

first = []
second = []
for a in sushi:
    if a:
        sort_a = sorted(a, reverse=True)
        flg = 1
        for num in sort_a:
            if flg:
                first.append((num, 1))
                flg = 0
            else:
                second.append((num, 0))
first.sort(reverse=True)                
second.sort(reverse=True)                

i = 0
j = 0
ans = 0

for i in range(len(first)):
    sum_first = 0
    sum_second = 0
    cnt = 0
    if i > K-1:
        break
    for j in range(i+1):
        sum_first += first[j][0]
        cnt += 1
    for k in range(K-i-1):
        if k > len(second) - 1:
            break
        sum_second += second[k][0]
  
    ans = max(ans, sum_first + sum_second + cnt*cnt)

print(ans)

Fixed submission:
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []

for nums in sushi.values():
    nums.sort(reverse=True)
    for i, j in enumerate(nums):
        if i == 0:
            first.append(j)
        else:
            second.append(j)
first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))

ans = 0

for i in range(K, 0, -1):
    neta = min(i, len(first))
    tmp = neta * neta + first_cum[neta]

    if K - neta > len(second):
        break
    
    tmp += second_cum[K - neta]

    ans = max(ans, tmp)

print(ans)