Buggy python submission (verdict: Time Limit Exceeded):
import itertools
import heapq
N, K = [int(x) for x in input().split()]
orig_td = [tuple([int(x) for x in input().split()]) for _ in range(N)]
td = sorted(orig_td.copy(), key=lambda x: x[1], reverse=True)
INF = 10**9
# print(td)


def score(a):
    kiso = sum([x[1] for x in a])
    bonus = len(set([x[0] for x in a]))**2
    return kiso + bonus


first = []
second = []
kind = set()
for i in range(K):
    if td[i][0] not in kind:
        first.append(td[i])
        kind.add(td[i][0])
    else:
        second.append(td[i])

best_ans = sum([x[1] for x in first + second]) + len(first)**2

for k in range(K, N):
    if len(second) == 0:
        break
    if td[k][0] in kind:
        continue

    a = second.pop(-1)
    kind.add(td[k][0])
    first.append(td[k])
    ans = sum([x[1] for x in first + second]) + len(first)**2
    # print(first, second, ans)
    if best_ans < ans:
        best_ans = ans
print(best_ans)


Fixed submission:
N, K = [int(x) for x in input().split()]
td = sorted([[int(x) for x in input().split()] for _ in range(N)],
            key=lambda x: x[1], reverse=True)
INF = 10**9
# print(td)

first = []
second = []
kind = set()
for i in range(K):
    if td[i][0] not in kind:
        first.append(td[i])
        kind.add(td[i][0])
    else:
        second.append(td[i])

yummy = sum([x[1] for x in first + second])
best_ans = yummy + len(kind)**2

for k in range(K, N):
    if len(second) == 0:
        break
    if td[k][0] in kind:
        continue

    a = second.pop(-1)
    yummy -= a[1]
    kind.add(td[k][0])
    first.append(td[k])
    yummy += td[k][1]
    ans = yummy + len(kind)**2
    # print(first, second, ans)
    if best_ans < ans:
        best_ans = ans
print(best_ans)

Buggy python submission (verdict: Time Limit Exceeded):
import itertools
N, K = [int(x) for x in input().split()]
orig_td = [tuple([int(x) for x in input().split()]) for _ in range(N)]
td = orig_td.copy()
INF = 10**9
# print(td)


def score(a):
    kiso = sum([x[1] for x in a])
    # print(a)
    bonus = len(set([x[0] for x in a]))**2
    return kiso + bonus


total = score(td)
for i in range(N, K, -1):
    # print("i", i)
    max_score = -INF
    useless_neta = 0
    for neta in td:
        te = td.copy()
        te.remove(neta)
        kari = score(te)
        if kari > max_score:
            max_score = kari
            useless_neta = neta
    td.remove(useless_neta)
    # print(td)
print(score(td))

# print("全列挙", max([score(x) for x in itertools.permutations(orig_td, K)]))


Fixed submission:
N, K = [int(x) for x in input().split()]
td = sorted([[int(x) for x in input().split()] for _ in range(N)],
            key=lambda x: x[1], reverse=True)
INF = 10**9
# print(td)

first = []
second = []
kind = set()
for i in range(K):
    if td[i][0] not in kind:
        first.append(td[i])
        kind.add(td[i][0])
    else:
        second.append(td[i])

yummy = sum([x[1] for x in first + second])
best_ans = yummy + len(kind)**2

for k in range(K, N):
    if len(second) == 0:
        break
    if td[k][0] in kind:
        continue

    a = second.pop(-1)
    yummy -= a[1]
    kind.add(td[k][0])
    first.append(td[k])
    yummy += td[k][1]
    ans = yummy + len(kind)**2
    # print(first, second, ans)
    if best_ans < ans:
        best_ans = ans
print(best_ans)