aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

abc 116 D - Various Sushi (python)

stack, dict, set 全部使った。盛りだくさん。
まず、降順ソートしてk個選んだ時に、L 種類の寿司が含まれているとき、当然それはL 種類の寿司を選ぶ美味しさの中での最大値となる。同様にして、k > L なら種類が被っている寿司を食べているので、k =L になるまで、被っていない中で最も美味しさが小さいものを、k+1個目以降でまだ食べてない種類の寿司が出てきたときに交換して、最大値が更新するかを調べたらよい。
被っているかは defaultdictで見て、k個までで被っているものは deque に入れておく。pop で出せば降順ソートしているので、被っているものの中での最小値が勝手に出てくる。

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from operator import itemgetter
from collections import defaultdict,deque

n,k = map(int,readline().split())
chk = []
for i in range(n):
  t,d = map(int,readline().split())
  chk.append((t,d))
chk = sorted(chk,key=itemgetter(1),reverse=True)

T = defaultdict(int)
d = deque()
s = set()
ans = []
aa = 0
for i in range(k):
  p,q = chk[i]
  T[p] += 1
  if T[p] > 1:
    d.append(q)
  s.add(p)
  aa += q
  
l = len(s)
aa += l**2
if k == n or l == k:
  print(aa)
  exit()
  
ans = aa
cnt = k-l
for p,q in chk[k:]:
  if T[p] == 0:
    T[p] += 1
    P = d.pop()
    aa -= P
    aa += q+2*l+1
    l += 1
    cnt -= 1
    ans = max(ans,aa)
  if cnt == 0:
    break
    
print(ans)