ABC116D

ABC116D の解説

atcoder.jp

解答コード

https://atcoder.jp/contests/abc116/submissions/31604534

問題

いくつかの種類に分かれているN個の寿司の中からK個の寿司を食べた時のコスト関数を最大化する問題です．

コスト = おいしさの総和 + (食べた種類)^2

解説

公式解説の想定解は，初めに美味しいK個の寿司を食べると仮定した貪欲法を試みますが，初めにすべての種類の寿司を食べると仮定して実装しました．こうすることで，種類ボーナスの管理が簡単になると思います．

概要図です．

ABC116D_overview

各種類1位にのネタをまず食べると仮定して配列A，
2位以下ネタをまとめて配列Bとします．

その後，Bの中にある寿司Bi（i=1, 2, ...）を食べた方が良いのかを判定していきます．
まだ食べる予定の寿司の個数がKに満たなければ，無条件でBiを食べることができます．

しかし，K個の寿司が食べる予定になっていれば，Aの中で最も不味いものを残ったBの中で最も美味しいものに入れ替えることで，種類ボーナスを上回る満足度が得られるかを貪欲に判定していきます．

あとは書くだけです．
概要図とソートの昇順降順が逆だったりしますが，そのままなので詳細は割愛します．

Parameters:
-----------
x : int
    種類の数です．
capa : int
    capacity のことです．あといくつ食べる余裕があるか．

# ABC116D

N, K = map(int, input().split())
T = [[] for _ in range(N)]
for _ in range(N):
    t, d = map(int, input().split())
    T[t-1].append(d)
A = []
B = []
for i in range(N-1, -1, -1):
    if T[i]:
        T[i].sort()
        A.append(T[i].pop())
        B.extend(T[i])
A.sort(reverse=True)
B.sort(reverse=True)
B.append(-1)
A = A[:K]
x = len(A)
capa = K - x
for i, b in enumerate(B):
    if capa:
        capa -= 1
    else:
        if A[-1]+x**2 < b+(x-1)**2:
            A.pop()
            x -= 1
        else:
            break
ans = sum(A) + sum(B[:i]) + x**2
print(ans)

以上，ABC116Dの解法でした．

＜参考解説リンク＞

公式: https://img.atcoder.jp/abc116/editorial.pdf