Buggy python submission (verdict: Runtime Error):
# 同じ種類だと、点数が高いやつを選ぶ
# 違う種類を選ぶと、ボーナス点が入るので、それをどうするか
# d  = sum(di) + x**x
# 種類xを固定すると？ O(log(N))でいける？
# ※※ 種類が増えるかどうかを降順で見ながらboolで置き換えるという発想。強い。

from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []
# まずは上からK個貪欲に取る
for i in range(k):
    d,t = sushi[i]
    heapq.heappush(sushiq,(d,t)) 
    sushidict[t] += 1
    cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(k,n):
    d,t = sushi[i]
    # すでにある種類の場合は大丈夫
    if sushidict[t] > 0 : continue
    # 種類が増える時、小さいものから入れ替えていく
    d_min,t_min = heapq.heappop(sushiq)
    # popして消えるようなものはNG
    while sushidict[t_min] == 1:
        d_min, t_min = heapq.heappop(sushiq)
    sushidict[t_min] -= 1
    sushidict[t] += 1
    n_kind += 1
    cnt  = cnt - d_min + d
    res.append(cnt+n_kind**2)
    if n_kind == n_kind_all:
        break

print(max(res))
# print(res)

Fixed submission:
"""
漸化式的に解く
一般性を失わない範囲をちゃんと考える。
"置き換え可能"というのをどの部分で持たせられるかを考える
今回の場合は、リスト作りの際に作ったので十分
"""

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] #変更できるもの(2個以上あるもの)



cnt = 0
for t,d in pairs[:K]: #K個上から順番に取ってみる
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

n = len(kinds)
cnt += len(kinds)**2
ans = cnt

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    # 交換できるものがなくなったら終わり
    if len(change) == 0:
        break
    # 同じものは交換しない
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)
# 集計
print(ans)

Buggy python submission (verdict: Runtime Error):
# 同じ種類だと、点数が高いやつを選ぶ
# 違う種類を選ぶと、ボーナス点が入るので、それをどうするか
# d  = sum(di) + x**x
# 種類xを固定すると？ O(log(N))でいける？
# ※※ 種類が増えるかどうかを降順で見ながらboolで置き換えるという発想。強い。

from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []
# まずは上からK個貪欲に取る
for i in range(k):
    heapq.heappush(sushiq,(d,t)) 
    d,t = sushi[i]
    sushidict[t] += 1
    cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(k,n):
    d,t = sushi[i]
    # すでにある種類の場合は大丈夫
    if sushidict[t] > 0 : continue
    # 種類が増える時、小さいものから入れ替えていく
    d_min,t_min = heapq.heappop(sushiq)
    # popして消えるようなものはダメ
    while sushidict[t_min] == 0:
        d_min, t_min = heapq.heappop(sushiq)
        print(d_min,t_min)
    sushidict[t_min] -= 1
    sushidict[t] += 1
    n_kind += 1
    cnt  = cnt - d_min + d
    res.append(cnt+n_kind**2)
    if n_kind == n_kind_all:
        break

print(max(res))

Fixed submission:
"""
漸化式的に解く
一般性を失わない範囲をちゃんと考える。
"置き換え可能"というのをどの部分で持たせられるかを考える
今回の場合は、リスト作りの際に作ったので十分
"""

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] #変更できるもの(2個以上あるもの)



cnt = 0
for t,d in pairs[:K]: #K個上から順番に取ってみる
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

n = len(kinds)
cnt += len(kinds)**2
ans = cnt

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    # 交換できるものがなくなったら終わり
    if len(change) == 0:
        break
    # 同じものは交換しない
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)
# 集計
print(ans)

Buggy python submission (verdict: Runtime Error):
# 漸化式的に解く
# 一般性を失わない範囲をちゃんと考える
# 2種類以上あるものは、同じ種類のものとして差し支えない

from collections import defaultdict
N, K = map(int,input().split())
dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]

# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = [0] * (N+1)

cnt = 0
oisii = []
for t,d in pairs[:K]: #K個上から順番に取ってみる
    kinds[t] += 1
    cnt += d
    # print(t,d)
    oisii.append([t,d])

point = sum([d for t,d in oisii]) + sum([1 for kind in kinds if kind!=0])**2
ans = point
for t,d in pairs[K:]:
    """
    置き換えようとしていたけど
    その必要はなく、
    ほしいのはmax値だけなので、配列そのものを更新し続ける必要はない
    """
    # if kinds[t] == 1:
    #     continue
    # else:
    #     index = K
    #     while True:
    #         t_min, d_min = oisii[index-1]
    #         # 2個以上ある種類のなかで最小値を求める
    #         if kinds[t_min] >= 2:
    #             break
    #         index -= 1
    #         if index == 0:
    #             break
    # if index == 0:
    #     break
    if kinds[t] >= 1:
        continue
    # とりあえず最小だったoisiiを置き換えていく
    t_min, d_min = oisii.pop()
    n = sum([1 for kind in kinds if kind!=0])
    # とりあえず種類を追加
    kinds[t] += 1
    kinds[t_min] -= 1
    point = point - (d_min + n**2)  + (d + (n+1)**2)
    # print(point)
    ans = max(point,ans)
# 集計
print(ans)


Fixed submission:
"""
漸化式的に解く
一般性を失わない範囲をちゃんと考える。
"置き換え可能"というのをどの部分で持たせられるかを考える
今回の場合は、リスト作りの際に作ったので十分
"""

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] #変更できるもの(2個以上あるもの)



cnt = 0
for t,d in pairs[:K]: #K個上から順番に取ってみる
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

n = len(kinds)
cnt += len(kinds)**2
ans = cnt

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    # 交換できるものがなくなったら終わり
    if len(change) == 0:
        break
    # 同じものは交換しない
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)
# 集計
print(ans)