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)