#AtCoder ABC116 -D- Python

マジでAtCoderの解法PDFを読むのが苦手なので(日本語力がないので)、お気持ちを書き残そうと思います。

atcoder.jp

まず最初に貪欲を考える癖があるので気を付けたいです。
今回の問題は種類ボーナスポイントさえなければもちろんおいしさを降順にソートして、k個足すだけで解くことができます。ですが、今回は基礎ポイントがあるのでこうは解けません。
しかし、美味しい順にk個食べても(わりと)いいポイントは得られそうです。ここから改善することを考えます。まず、おいしさ順に上からk個食べた寿司の集合をSとします。よりポイントを高めるためにはSに含まれる種類が増えなければなりません。(もし増えなければ今のSのスコアが最大のはずです) ということでSの中で同じ種類で2つ以上食った寿司の中から一番まずいもの(bと名付けます)を吐きだします。そして食べたことのない一番おいしい寿司(gと名付けます)を食べます。つまり、Sからbを消してaを追加します。しかしこれでスコアが増えるかどうかはわかりません。よって、この作業を繰り返してSのスコアのmaxを取るのがいいでしょう。

コードに起こすと以下のようになります。

n, k = map(int, input().split())
sushi = []
for i in range(n):
    tmp = list(map(int, input().split()))
    sushi.append(tmp)
eaten = set() # 食べた種類を記憶(集合型)
chohuku = [] # 同じ種類で二回以上食べた寿司のスコアを記録
scores = [] # Sのスコアのリスト

sushi.sort(key=lambda x:x[1])
sushi.reverse()

tmp = 0 # Sのおいしさの合計を保持
for i in range(k): 
    if sushi[i][0] in eaten: # 同じ種類の寿司を食べていた場合
        chohuku.append(sushi[i][1]) 
    else:
        eaten.add(sushi[i][0]) 
    tmp += sushi[i][1] 
chohuku.sort() 
scores.append(tmp+len(eaten)**2) 

for i in range(k,n): # 残りの寿司から美味しい順に食べるものを選ぶ
    if not sushi[i][0] in eaten and chohuku != []: # 重複が存在し、また食べていない種類のとき
        tmp -= chohuku.pop(0) 
        eaten.add(sushi[i][0]) 
        tmp += sushi[i][1] 
        scores.append(tmp+len(eaten)**2)
        
print(max(scores))