Buggy python submission (verdict: Time Limit Exceeded):
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

bucket = [0]*(N+1)
for t in n_dm_set:
    bucket[t] = -1

d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    t = t_arr[i]
    if bucket[t] != 0:
        continue
    d = d_arr[i]
    bucket[t] = d
    if len(small_arr) == 0:
        break
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)


Fixed submission:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)

Buggy python submission (verdict: Time Limit Exceeded):
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

bucket = [0]*(N+1)
for t in n_dm_set:
    bucket[t] = -1

d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    t = t_arr[i]
    if bucket[t] != 0:
        continue
    d = d_arr[i]
    bucket[t] = d
    if len(small_arr) == 0:
        break
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    nt_x += 1
    candidate_dict[nt_x] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)


Fixed submission:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)

Buggy python submission (verdict: Time Limit Exceeded):
N,K = (int(x) for x in input().split())
td_arr = []
for _ in range(N):
    td = tuple(int(x) for x in input().split())
    td_arr.append(td)

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

n_dm_set = set()
small_arr = []
d_dm = 0
for t,d in td_arr[:K]:
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

prev_t = -1
other_arr = []
for t,d in sorted(td_arr, key=lambda x:x[0]):
    if t == prev_t:
        continue
    if t in n_dm_set:
        continue
    prev_t = t
    other_arr.append((t,d))
other_arr.sort(key=lambda x:x[1], reverse=True)

d_x = d_dm
nt_x = len(n_dm_set)
for t,d in other_arr:
    if len(small_arr) == 0:
        break
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    nt_x += 1
    candidate_dict[nt_x] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)


Fixed submission:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)

Buggy python submission (verdict: Time Limit Exceeded):
N,K = (int(x) for x in input().split())
td_arr = []
for _ in range(N):
    td = tuple(int(x) for x in input().split())
    td_arr.append(td)

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

n_dm_set = set()
small_arr = []
d_dm = 0
for t,d in td_arr[:K]:
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
for t,d in td_arr[K:]:
    if t in n_x_set:
        continue
    if len(small_arr) == 0:
        break
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    candidate_dict[len(n_x_set)] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)


Fixed submission:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)

Buggy python submission (verdict: Runtime Error):
N,K = (int(x) for x in input().split())
td_arr = []
for _ in range(N):
    td = tuple(int(x) for x in input().split())
    td_arr.append(td)

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

n_dm_set = set()
small_arr = []
d_dm = 0
for t,d in td_arr[:K]:
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
for t,d in td_arr[K:]:
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    candidate_dict[len(n_x_set)] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)


Fixed submission:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)