Problem: There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i. You are choosing K among these N pieces to eat. Your "satisfaction" here will be calculated as follows:    * The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".   * The base total deliciousness is the sum of the deliciousness of the pieces you eat.   * The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.  You want to have as much satisfaction as possible. Find this maximum satisfaction.

Submission attempts (in order):


# d  = sum(di) + x**x



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 = []

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)
    
    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))






from collections import defaultdict
N, K = map(int,input().split())

pairs = [list(map(int,input().split())) for _ in range(N)]

pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] 



cnt = 0
for t,d in pairs[:K]: 
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

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


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)




# d  = sum(di) + x**x



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 = []

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)
    
    while sushidict[t_min] == 1 and sushiq:
        d_min, t_min = heapq.heappop(sushiq)
    
    if not sushiq:
        break
    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)

Verdicts: Runtime Error

Wrong Answer

Accepted