TnT/Multi_CodeNet4Repair huggingface dataset:
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):
N,K=map(int,input().split())

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0

number=set()
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.append(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[i][0] not in number:
    number.append(td[i][0])
    tempanswer=tempanswer-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)


N,K=map(int,input().split())

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0
number=set()
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.add(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[index][0] not in number:
    number.add(td[index][0])
    tempanswer=tempanswer+td[index][1]-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)

Verdicts: Runtime Error

Accepted
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):
import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()


import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            s.add(t)
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()


Verdicts: Wrong Answer

Accepted
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):
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r = []
    variety = set([])
    ans = 0
    for _c in choice:
        ans += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    ans += len(variety)**2
    i = k
    while len(variety) < min(max_variety, k):
        while td[i][0] in variety:
            i += 1
        variety.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))


def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    r, v, b = [], set([]), 0
    for _c in choice:
        b += _c[1]
        if _c[0] in v:
            r.append(_c[1])
        else:
            v.add(_c[0])
    b += len(v)**2
    i = k
    ans = [b]
    while len(v) < min(max_variety, k):
        while td[i][0] in v:
            i += 1
        v.add(td[i][0])
        ans.append(ans[-1] - r.pop() + td[i][1] + 2 * len(v) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))


Verdicts: Runtime Error

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

# -*-coding: utf-8 -*-

import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [0], [0]
    st = set()

    for t, d in S:
        if t in st:
            
            y_0.append(y_0[-1]+d)
        else:
            st.add(t)
            y_1.append(y_1[-1]+d)
    

    max_sum = -1

    for i in range(1,k+1):
        print(i)
        if k-i >= len(y_0):
            continue
        if i >= len(y_1):
            break
        max_sum = max(max_sum, y_0[k-i]+y_1[i]+i**2)
    
    print(max_sum)
    return


if __name__ == '__main__':
    main()


# -*-coding: utf-8 -*-

import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [0], [0]
    st = set()

    for t, d in S:
        if t in st:
            
            y_0.append(y_0[-1]+d)
        else:
            st.add(t)
            y_1.append(y_1[-1]+d)
    

    max_sum = -1

    for i in range(1,k+1):
        if k-i >= len(y_0):
            continue
        if i >= len(y_1):
            break
        max_sum = max(max_sum, y_0[k-i]+y_1[i]+i**2)
    
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Verdicts: Wrong Answer

Accepted
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):
n, k = map(int, input().split())
TD = [[] for i in range(n)]
for i in range(n):
  TD[i] = list(map(int, input().split()))
  
TD.sort(key = lambda x:x[1], reverse = True)

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2
huku.sort(key = lambda x:x[1], reverse = True)

ansl = []
ansl.append(ans)
for i,j in TD[k:]:
  if i in su:
    continue
  elif:
    huku == []
    break
  else:
    temp = ansl[-1] + j - huku.pop()[1] + (len(su)+1)**2 -len(su)**2
    su.add(i)
    ansl.append(temp)

print(max(ansl))

n, k = map(int, input().split())
TD = [[] for i in range(n)]
for i in range(n):
  TD[i] = list(map(int, input().split()))
  
TD.sort(key = lambda x:x[1], reverse = True)

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2
huku.sort(key = lambda x:x[1], reverse = True)

ansl = []
ansl.append(ans)
for i,j in TD[k:]:
  if i in su:
    continue
  elif huku == []:
    break
  else:
    temp = ansl[-1] + j - huku.pop()[1] + (len(su)+1)**2 -len(su)**2
    su.add(i)
    ansl.append(temp)

print(max(ansl))


Verdicts: Runtime Error

Accepted
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):
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
import heapq
n,k=IL()
nl=dict()
for i in range(n+1):
    nl[i]=0
l=[]
ans=[]
f=[]
el=[]
heapq.heapify(el)
s=0
for i in range(n):
    t,d=IL()
    l.append([d,t])
l.sort(reverse=True)
for i in range(k):
    t=l.pop(0)
    el.heappush(t)
    nl[t[1]]+=1
    s+=t[0]
    if t[1] not in f:
        f.append(t[1])
u=len(f)
ans.append(s+u*u)
for i in range(u,k+1):
    if len(l)==0:
        break
    while True:
        t=el.heappop(0)
        if nl[t[1]]!=1:
            nl[t[1]]-=1
            s-=nl[t[0]]
            break
    while len(l)>0:
        t=l.pop(0)
        if nl[t[1]]==0:
            nl[t[1]]+=1
            s+=t[0]
            el.heappush(t)
            break
    ans.append(s+(u+i+1)**2)
    if len(l)==0:
        break
print(max(ans))

def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
n,k=IL()
top=dict()
sub=[]
for i in range(n):
    t,d=IL()
    if t in top:
        sub.append(min(top[t],d))
        top[t]=max(top[t],d)
    else:
        top[t]=d

top=sorted(top.items(),key=lambda x:-x[1])
sub.sort(reverse=True)

rtop=[top[0][1]]
rsub=[0]
for i in range(1,len(top)):
    rtop.append(rtop[-1]+top[i][1])
for i in range(len(sub)):
    rsub.append(rsub[-1]+sub[i])

ans=0
for i in range(k):
    if i<len(rtop) and 0<=k-i-1<len(rsub):
        ans=max(ans,rtop[i]+rsub[k-i-1]+(i+1)**2)
print(ans)

Verdicts: Runtime Error

Accepted
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):
N, K = map(int, input().split())
sushi = []
for i in N:
    s = list(map(int, input().split()))
    sushi.append(s)


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

def find_next(start, curr, point, tlist, tcount):
    if start == N:
        return [point, curr]
    max_c = curr
    max_p = point
    for s in sushi[start:]:
        if tlist[s[0]] > 0:
            continue
        for ri in range(len(curr) - 1, -1, -1):
            r = curr[ri]
            if tlist[r[0]] >= 2:
                c2 = [c for c in curr]
                t2 = [t for t in tlist]
                del c2[ri]
                c2.append(s)
                t2[r[0]] -= 1
                t2[s[0]] = 1
                p2 = point - r[1] + s[1] + (tcount + 1) * (tcount + 1) - tcount * tcount
                np, pc = find_next(start + 1, c2, p2, t2, tcount + 1)

                if np > max_p:
                    max_p = np
                    max_c = pc
                break
    return [max_p, max_c]

result = sushi[0: K]
type_count = [0] * (N + 1)

for r in result:
    type_count[r[0]] += 1
type_num = 0
for t in type_count:
    if t > 0:
        type_num += 1
point = type_num * type_num
for r in result:
    point += r[1]
ans_p, ans_c = find_next(K, result, point, type_count, type_num)

print(ans_p)


N, K = map(int, input().split())
sushi = []
for i in range(N):
    s = list(map(int, input().split()))
    sushi.append(s)


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

# pick from first
point = 0
neta_set = set()
removable_idx = list()
for i in range(K):
    s = sushi[i]
    point += s[1]
    if s[0] not in neta_set:
        neta_set.add(s[0])
    else:
        removable_idx.append(i)
point += len(neta_set) **2
max_point = point

# pick latest
for s in sushi[K:]:
    if len(removable_idx) == 0:
        break
    if s[0] in neta_set:
        continue
    idx = removable_idx[-1]
    count = len(neta_set)
    point += s[1] - sushi[idx][1]
    point += (count + 1) ** 2 - count ** 2
    neta_set.add(s[0])
    removable_idx = removable_idx[:-1]

    if point > max_point:
        max_point = point
print(max_point)

Verdicts: Runtime Error

Accepted
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):
nklist = list(map(int,input().split()))
n, k= nklist[0], nklist[1]

lst = [tuple(map(int, input().split())) for i in range(n)]

max_pt = 0
for comb in itertools.combinations(lst, k):
    kinds_set = set()
    umami = 0
    for tpl in comb:
        kinds_set.add(tpl[0])
        umami += tpl[1]
    ttl_pt = len(kinds_set)**2 + umami
    #print('ttl_pt:', ttl_pt,' max_pt:', max_pt)
    if ttl_pt > max_pt:
        max_pt = ttl_pt
print(max_pt)

n, k = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x:-x[1])
 
ans = 0
s = set()
q = []
 
for i in range(k):
    ans += a[i][1]
    if a[i][0] in s:
        q.append(a[i][1])
    s.add(a[i][0])
 
ans += len(s)**2
k1 = k
ans1 = ans
 

while k1 < n and q:
    if a[k1][0] not in s:
        ans1 = ans1+a[k1][1]-q.pop()-len(s)**2+(len(s)+1)**2
        ans = max(ans, ans1)
        s.add(a[k1][0])
    k1 += 1
 
print(ans)

Verdicts: Runtime Error

Accepted
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):
from itertools import groupby
from heapq import heapify, heappop

N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]


def point(sushi):
    t, d = zip(*sushi)
    return sum(d) + len(set(t))**2


srt = sorted(sushi, key=lambda x: x[1], reverse=True)
eat, non_eat = srt[:K], srt[K:]

maxs = []
rems = []
kinds = set()
for k, g in groupby(sorted(eat, key=lambda x: x[0]), key=lambda x: x[0]):
    _, d = zip(*g)
    maxs.append(d.pop(0))
    rems.extend(d)
    kinds.add(k)
heapify(rems)
ans = point(eat)

for k, g in groupby(sorted(non_eat, key=lambda x: x[0]), key=lambda x: x[0]):
    if k in kinds:
        continue
    _, d = zip(*g)
    maxs.append(d.pop(0))
    heappop(rems)
    kinds.add(k)
    ans = max(ans, point(maxs + rems))

print(ans)

from heapq import heappop, heappush, heapify

N, K = map(int, input().split())
sushi = [(-d, t) for t, d in map(int, input().split())]
heapify(sushi)

maxs = []
rems = []
kinds = set()
for _ in range(K):
    d, t = heappop(sushi)
    if t in kinds:
        heappush(rems, -d)
    else:
        maxs.append(-d)
        kinds.add(t)
ans = sum(maxs) + sum(rems) + len(kinds)**2

while True:
    if not sushi or not rems:
        break
    d, t = heappop(sushi)
    if t in kinds:
        continue
    else:
        heappop(rems)
        maxs.append(-d)
        kinds.add(t)
        ans = max(ans, sum(maxs) + sum(rems) + len(kinds)**2)

print(ans)

N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi.sort(key=lambda x: x[1], reverse=True)

score = 0

remains = []
kinds = set()
for t, d in sushi[:K]:
    if t in kinds:
        remains.append(d)
    kinds.add(t)
    score += d
score += len(kinds)**2

test = score
for t, d in sushi[K:]:
    if not remains:
        break
    if t in kinds:
        continue
    kinds.add(t)
    rm = remains.pop()
    test = test + (d - rm) + (len(kinds)*2 - 1)
    score = max(score, test)

print(score)

Verdicts: Runtime Error

Runtime Error

Accepted
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):
N,K=map(int,input().split())

dtlist = []
tdlist = []
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tdlist.append((t,d))
  
dtlist.sort(reverse=True)
tdlist.sort(reverse=True)
print(dtlist)
print(tdlist)
  
sum_d=0
set_t=set()
for i in range(K):
  d,t=dtlist[i]
  sum_d += d
  set_t.add(t)
  
print(sum_d+len(set_t)*len(set_t))

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

dtlist=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tset_all.add(t)
dtlist.sort()
#print(heapq)

dtlist_Krev=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  d,t=dtlist.pop()
  dtlist_Krev.append((d,t))
  dsum+=d
  tdic[t]+=1
#dtlist_Krev.sort(reverse=True)

t0=len(tdic)
max_answer=dsum+(t0**2)
#print(t0,max_answer)

for i in range(t0+1,min(K,len(tset_all))+1):
  while(True):    
    d1,t1=dtlist_Krev.pop()
    if tdic[t1]>1:
      tdic[t1]-=1
      dsum-=d1
      break
      
  while(True):        
    d2,t2=dtlist.pop()
    if tdic[t2]==0:
      tdic[t2]=1
      dsum+=d2      
      break
          
  answer_i=dsum+(i**2)
  #print(i,answer_i)
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

Verdicts: Wrong Answer

Accepted
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):
n,k=map(int,input().split())
X=[]
for i in range(n):
    t,d=map(int,input().split())
    X.append([t,d])
    
X.sort(key=lambda x:x[1],reverse=True)

V=set()
f=0
S=[]

for t,d in X[:k]:
    f+=d
    if t in V:
        S.append(t)
    else:
        V.add(t)

point=f+len(V)**2

ans=point

for t,d in X[k:]:
    if not S:
        break
    if not t in V:
        point+=d-S.pop()+1+len(V)*2
        V.add(t)
        ans=max(ans,point)
print(ans)
    

n,k=map(int,input().split())
X=[]
for i in range(n):
    t,d=map(int,input().split())
    X.append([t,d])
    
X.sort(key=lambda x:x[1],reverse=True)

V=set()
point=0
S=[]

for t,d in X[:k]:
    point+=d
    if t in V:
        S.append(d)
    else:
        V.add(t)
point+=len(V)**2

ans=point

for t,d in X[k:]:
    if not S:
        break
    if not t in V:
        point+=d-S.pop()+1+len(V)*2
        V.add(t)
        ans=max(ans,point)

print(ans)

Verdicts: Wrong Answer

Accepted
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):
from collections import defaultdict
N, K = map(int, input().split())
T = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    T[t].append(d)
S = []
U = []
for k, v in T.items():
    v.sort(reverse=1)
    U.append((v[0], k))
    for w in v[1:]:
        S.append((w, k))
S.sort(reverse=1)
U.sort(reverse=1)
D = defaultdict(int)
cnt = 0
D[U[0][1]] += 1
t = 1
cnt += 1
res = U[0][0]
S = S[:K-1]
for w, k in S:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += w
ans = 0
for v, k in U[1:]:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += v
    while cnt > K:
        cnt -= 1
        w, l = S.pop()
        D[l] -= 1
        if D[l] == 0:
            t -= 1
        res -= w
    print(t, res, cnt)
    ans = max(ans, t*t + res)
print(ans)

from collections import defaultdict
N, K = map(int, input().split())
T = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    T[t].append(d)
S = []; U = []
for k, v in T.items():
    v.sort(reverse=1)
    U.append((v[0], k))
    for w in v[1:]:
        S.append((w, k))
S.sort(reverse=1)
U.sort(reverse=1)
D = defaultdict(int)
t = cnt = 0
res = 0
for w, k in S:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += w
ans = 0
can = 1
for v, k in U:
    D[k] += 1
    if D[k] == 1:
        t += 1
    cnt += 1
    res += v
    while cnt > K:
        if not S:
            can = 0
            break
        cnt -= 1
        w, l = S.pop()
        D[l] -= 1
        if D[l] == 0:
            t -= 1
        res -= w
    if not can:
        break
    ans = max(ans, t*t + res)
print(ans)

Verdicts: Runtime Error

Accepted
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):
from operator import itemgetter

n, k = map(int, input().split(" "))
sushis = []
flags = [False for i in range(n+10)]
sushis2 = []
sushis22 = []
sum2 = 0
count2 = 0
for i in range(n):
    t, d = (map(int, input().split(" ")))
    sushis.append([t, d])

    if not flags[t]:
        sushis22.append(d)
        flags[t] = True
    else:
        sushis2.append(d)
# calc variety
if len(sushis22) <= k:
    count2 += len(sushis22)
    sum2 += sum(sushis22)
    sum2 += len(sushis22) * len(sushis22)
else:
    sum2 += sum(sorted(sushis22, reverse=True)[0:k])
    sum2 += k*k
    count2 += k

if count2 < k:
    sushis2 = sorted(sushis2, reverse=True)
    sum2 += sum(sushis2[0:k-count2])

# calc score
sum1 = 0
set1 = []
count = k
temp_score = 0
temp_i = 0
sushis = sorted(sushis, key=itemgetter(1), reverse=True)

temp_set = []
for i, sushi in enumerate(sushis):
    sum1 += sushi[1]
    if count == 1:
        temp_set.append(sushi[0])
        temp_i = i
        temp_score = sushi[1]
        break
    group = sushi[0]
    set1.append(group)
    count -= 1


set1 = set(set1)

while temp_i+1 < n:
    if temp_score == sushis[temp_i+1]:
        temp_set.append(sushis[temp_i+1][0])
        temp_i += 1
    else:
        break

if len(set1 | set(temp_set)) != len(set1):
    sum1 += (len(set1)+1)*(len(set1)+1)
else:
    sum1 += len(set1) * len(set1)
print(sum1)
print(sum2)

print(max(sum1, sum2))


from operator import itemgetter

n, k = map(int, input().split(" "))
sushis = [list(map(int, input().split(" "))) for i in range(n)]
flags = [False for i in range(n+10)]

# calc score
sushis = sorted(sushis, key=itemgetter(1), reverse=True)
duplicated_set = []
variety = 0
eat = 0
max_score = 0
sum1 = 0
for i, sushi in enumerate(sushis):
    t, d = sushi
    if eat < k:
        sum1 += d
        eat += 1
        if flags[t]:
            duplicated_set.append(d)
        else:
            variety += 1
            flags[t] = True
        if eat == k:
            sum1 += variety * variety
            max_score = sum1
            duplicated_set = sorted(duplicated_set, reverse=True)
    else:
        if variety == k or len(duplicated_set) == 0:
            break
        if flags[t]:
            pass
        else:
            variety += 1
            flags[t] = True
            sum1 = sum1 + d - duplicated_set.pop() + 2*(variety) - 1
            max_score = max(max_score, sum1)

print(max_score)


Verdicts: Wrong Answer

Accepted
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):
n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = set()
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = sum(ans) + sum(aru) + s**2

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans.append(nai[i][1])
        s += 1
    m = max(m, sum(ans) + sum(aru) + (s + i)**2)
print(m)


n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = set()
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = sum(ans) + sum(aru) + s**2
print(ans, aru, s)

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans.append(nai[i][1])
        s += 1
        print(ans, aru, s)
    m = max(m, sum(ans) + sum(aru) + s**2)
print(m)


n, k = map(int, input().split())
nai = [list(map(int, input().split())) for _ in range(n)]

nai = sorted(nai, key=lambda x: -x[1])

ans = 0
aru = []
spe = set()
m = 0
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans+=anstmp[1]
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])
sumaru = sum(aru)
m = ans + sumaru + len(spe)**2

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        sumaru-=aru.pop()
        ans+=nai[i][1]
        spe.add(nai[i][0])
    m = max(m, ans + sumaru + len(spe)**2)
print(m)


Verdicts: Wrong Answer

Runtime Error

Accepted
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):
#import collections
n, k = map(int, input().split())
c = []
for i in range(n):
    c.append([int(m) for m in input().split()])
c.sort(key=lambda x:(x[0], x[1]))
f = sorted(c, key=lambda x:(x[1], x[0]), reverse=True)
donyoku = 0
donset = set()
for i in range(k):
    donyoku += f[i][1]
    donset.add(f[i][0])
donyoku += len(donset)**2

t = []
d = []
for i in range(n):
    pi, qi = c[i][0], c[i][1]
    t.append(pi)
    d.append(qi)
spe = len(set(t))
stili = []
real = 0
reaset = set()
ws = set()
i = 0
p = [0]*n
while True:
    if f[i][0] not in ws:
        real += f[i][1]
        ws.add(f[i][0])
        p[i] = 1
        i += 1
    if i == len(donset):
        break
kosu = len(donset)
yy = real
u = 0
for i in range(len(donset) + 1, k + 1):
    while True:
        if f[i][0] not in ws:
            real += f[i][1]
            ws.add(f[i][0])
            kosu += 1
            break
        else:
            i += 1
        if i == n:
            u = 1
            break
    if u == 1:
        break
    q = 0
    while True:
        if q < n:
            if p[q] == 0:
                real += f[q][1]
                kosu += 1
                if kosu == k:
                    break
        else:
            break
        q += 1
    real += len(ws) ** 2
    stili.append(real)
    real = yy


if stili != []:
    rr = min(stili)
    print(min(donyoku,rr))
else:
    print(donyoku)

from heapq import heappush, heappop
#import collections
n, k = map(int, input().split())
c = []
for i in range(n):
    c.append([int(m) for m in input().split()])
c.sort(key=lambda x:(x[0], x[1]))
f = sorted(c, key=lambda x:(x[1], x[0]), reverse=True)
donyoku = 0
donset = set()

notfirst =[]
for i in range(k):
    if f[i][0] not in donset:
        donset.add(f[i][0])
    else:
        heappush(notfirst, f[i][1])
    donyoku += f[i][1]
donyoku += len(donset)**2

candidate = []
candidateset = set()
donyokulength = len(donset)
for i in range(k, n):
    if f[i][0] not in donset:#
        candidate.append(-f[i][1])
        candidateset.add(f[i][0])
        donset.add(f[i][0])

satili = [donyoku]
donyoku -= donyokulength**2
for i in range(len(candidateset)):
    if notfirst == []:
        break
    else:
        donyoku -= heappop(notfirst)
        donyoku += -heappop(candidate)
        donyoku += (donyokulength + i + 1) ** 2
        satili.append(donyoku)
    donyoku -= (donyokulength + i + 1) ** 2




print(max(satili))


Verdicts: Wrong Answer

Accepted
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):
# -*- coding utf-8 -*-

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key = lambda x: x[1], reverse = True)    # sort by d

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()

for i in range(k):
    op += sushi[i][1]

    t = sushi[i][0]
    if t_taberu.isdisjoint({t}):    # new type
        t_taberu.add(t)
    else:
        pop_cands.append(i)

    taberu_ind.add(i)
mp_max = op + len(t_taberu)**2


i += 1
j = i
for i in range(i, n):
    pop_ind = pop_cands.pop()
    taberu_ind.remove(pop_ind)

    for j in range(j, n):
        t = sushi[j][0]
        if t_taberu.isdisjoint({t}):
            t_taberu.add(t)
            break
    else:
        break
    
    # calc points
    op -= sushi[pop_ind][1]
    op += sushi[j][1]
    mp = op + len(t_taberu)**2
    mp_max = max(mp_max, mp)

    taberu_ind.add(i)

print(mp_max)



# -*- coding utf-8 -*-

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key = lambda x: x[1], reverse = True)    # sort by d

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()

for i in range(k):
    op += sushi[i][1]

    t = sushi[i][0]
    if t_taberu.isdisjoint({t}):    # new type
        t_taberu.add(t)
    else:
        pop_cands.append(i)

    taberu_ind.add(i)
mp_max = op + len(t_taberu)**2


i += 1
j = i
for i in range(i, n):
    if len(pop_cands) == 0:
        break

    pop_ind = pop_cands.pop()
    taberu_ind.remove(pop_ind)

    for j in range(j, n):
        t = sushi[j][0]
        if t_taberu.isdisjoint({t}):
            t_taberu.add(t)
            break
    else:
        break
    
    # calc points
    op -= sushi[pop_ind][1]
    op += sushi[j][1]
    mp = op + len(t_taberu)**2
    mp_max = max(mp_max, mp)

    taberu_ind.add(i)    # not necessary actually
print(mp_max)



Verdicts: Runtime Error

Accepted
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):
from collections import defaultdict

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = defaultdict(int)
x = list(range(n - 1, n - k - 1, -1))
z = [0]
for i in range(n - k, n):
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]

y = []
yd = list(t.keys())
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])
y = y[::-1]
k = 1

for k in range(1, len(y) + 1):
    for i in range(len(x) - 1, -1, -1):
        xi = x[i]
        if t[sushi[xi][0]] > 1:
            t[sushi[xi][0]] -= 1
            z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])
            x.pop(i)
            break

p = len(t.keys())
ans = 0
for i in range(len(z)):
    ans = max(z[i] + (p + i) ** 2)
print(ans)


n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        x.append(di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
idx = 1
for i in range(n - k - 1, -1, -1):
    if idx == len(x) + 1:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - x[-idx]
        p += 1
        idx += 1
        ans = max(ans, z + p ** 2)

print(ans)


Verdicts: Runtime Error

Accepted
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):
from heapq import*
n,k=map(int,input().split())
l=[tuple(map(int,input().split()))for i in range(n)]
heapify(sorted(l,key=lambda x:x[1])[::-1])
print(l[:k])
print(l[k:])


from collections import*
n,k=map(int,input().split())
sushi=sorted([tuple(map(int,input().split()))for i in range(n)],key=lambda x:x[1],reverse=1)
first=defaultdict(list)
second=defaultdict(list)
menu=set()

for _,v in sushi[:k]:
    menu.add(_)
    first[_].append(v)
second=sushi[k:][::-1]
score=sum(map(sum,first.values()))
change=[]
for _,v in first.items():
    for i in v[1:]:
        change.append((_,i))
change.sort(key=lambda x:x[1],reverse=1)

#print(first,change,second)
ans=[score+len(first)**2]
while change and second:
    _,s=change.pop()
    score-=s
    try:
        while second[-1][0] in menu:
            pp,ss=second.pop()
        pp,ss=second.pop()
        menu.add(pp)
        score+=ss
        ans.append(score+len(menu)**2)
    except:
        break
print(max(ans))



Verdicts: Wrong Answer

Accepted
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):
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))

A.sort(reverse=True)
X = []
Y = []
st = set()

print(A)

for i in range(k):
    d,t = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)

print(X)

x = sum(X)
y = sum(Y)
v = len(Y)
ans = x + y + v**2
for i in range(k, n):
    if not X:
        break
    d, t = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v**2)

print(ans)


n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))

A.sort(reverse=True)
X = []
Y = []
st = set()

for i in range(k):
    d,t = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)

x = sum(X)
y = sum(Y)
v = len(Y)
ans = x + y + v**2
for i in range(k, n):
    if not X:
        break
    d, t = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v**2)

print(ans)


Verdicts: Wrong Answer

Accepted
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):
N,K,*L = map(int, open(0).read().split())
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  ls.append((d,t))
ls.sort()
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = ls.pop()
  ans += d
  if t not in S:
    S.add(t)
  else:
    pre.append(d)
pre.sort(reverse=True)
ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = ls.pop()
  if t in S:
    continue
  S.add(t)
  m -=pre.pop()
  m += d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

N,K,*L = map(int, open(0).read().split())
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  ls.append((d,t))
ls.sort()
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = ls.pop()
  ans += d
  if t not in S:
    num += 1
    S.add(t)
  else:
    pre.append(d)
pre.sort(reverse=True)
ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = ls.pop()
  if t in S:
    continue
  S.add(t)
  m -=pre.pop()
  m += d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

Verdicts: Wrong Answer

Accepted
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):
import numpy as np

n, k = map(int, input().split())
mx = {}
for _ in range(n):
    t, d  = map(int, input().split())
    if t not in mx:
        mx[t] = d
    elif mx[t] < d:
        mx[t] = d
    sushi.append((t,d))


taste_sorted_sushi = sorted(sushi, key=lambda x:x[1], reverse=True)


selected_sushi = taste_sorted_sushi[:k]



S = set()
q = []
for s in selected_sushi:
    if s[0] in S:
        q.append(s)
    else:
        S.add(s[0])

q = sorted(q, key=lambda x:x[1], reverse=True)
ans = sum(selected_sushi[:][1])+len(S)**2

S = set(selected_sushi[:][0])
for i in range(k,n):
    if taste_sorted_sushi[k] not in S:
        ans = max(ans, ans-q.pop()+taste_sorted_sushi[k]-len(S)**2 + (len(S)+1)**2
        S.add(taste_sorted_sushi[k][0])

print(ans)

import numpy as np

n, k = map(int, input().split())
mx = {}
sushi = []
for _ in range(n):
    t, d  = map(int, input().split())
    if t not in mx:
        mx[t] = d
    elif mx[t] < d:
        mx[t] = d
    sushi.append([t,d])


taste_sorted_sushi = sorted(sushi, key=lambda x:x[1], reverse=True)


selected_sushi = np.array(taste_sorted_sushi[:k], dtype=np.int64)



S = set()
q = []
for s in selected_sushi:
    if s[0] in S:
        q.append(s[1])
    else:
        S.add(s[0])

q = sorted(q, reverse=True)
ans = sum(selected_sushi[:,1])+len(S)**2
tmp = ans

#print(selected_sushi, selected_sushi[:,1], 'ans', ans, S, q)

#print(S)
for i in range(k,n):
    if not q:
        break
    if taste_sorted_sushi[i][0] not in S:
        #print('ans',ans)
        
        tmp = tmp-q.pop()+taste_sorted_sushi[i][1]-len(S)**2 + (len(S)+1)**2
        ans = max(ans, tmp)
        S.add(taste_sorted_sushi[i][0])

print(ans)

Verdicts: Runtime Error

Accepted
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):
N,K = list(map(int,input().split()))
TDlist = []

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

TDlist.reverse(key = lambda x: x[1])
kinds = []
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.append(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer
max_score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.append(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    max_score = max(score,answer)
    

answer = max(answer,max_score)

print(answer)

N,K = list(map(int,input().split()))
TDlist = []

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

TDlist= sorted(TDlist, key = lambda x: x[1],reverse=True)
kinds = set()
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.add(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.add(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    answer = max(score,answer)

print(answer)

Verdicts: Runtime Error

Accepted