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 heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()

for i in range(len(TD)):
    flg = 0
    if TD[i][0] not in muptiple:
        muptiple.add(TD[i][0])
        flg = 1
    TD[i][0] = flg
TD.reverse()

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = TD.pop()
    vari += v[0]
    point += v[1]
    if not v[0]:
        heappush(ret, t)

total = vari ** 2 + point

while len(TD) > 0 and len(ret) > 0:
    v = TD.pop()
    if v[0] == 0:
        continue
    pre_t = heappop(ret)
    point -= pre_t - v[1]
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)

n, k = map(int,input().split())
inputarray = []
for i in range(n):
    inputarray.append([int(i) for i in input().split()])
t = [i[0] for i in inputarray]
d = [i[1] for i in inputarray]
print(n,k)
print(t)
print(d)
def combinationListRecursive(data, r):
  if r == 0 or r > len(data):
    return []

  result = []
  _combinationListRecursive(data, r, 0, [], result)
  return result


def _combinationListRecursive(data, r, start, progress, result):
  if r == 0:
    result.append(progress)
    return

  for i in range(start, len(data)):
    _combinationListRecursive(data, r - 1, i + 1, progress + [data[i]], result)

result = combinationListRecursive(range(len(d)), k)
best = 0
for i in result:
    tarray = []
    dtmp = 0
    for k in range(k):
        dtmp += d[i[k]]
        tarray.append(t[i[k]])
    tlist = list(set(tarray))
    ttmp = len(tlist) ** 2
    if best < dtmp + ttmp:
        best = dtmp + ttmp
print(best)

n, k = map(int,input().split())
t = [int(input().split()[0]) for i in range(n)]
d = [int(input().split()[1]) for i in range(n)]
def combinationListRecursive(data, r):
  if r == 0 or r > len(data):
    return []

  result = []
  _combinationListRecursive(data, r, 0, [], result)
  return result


def _combinationListRecursive(data, r, start, progress, result):
  if r == 0:
    result.append(progress)
    return

  for i in range(start, len(data)):
    _combinationListRecursive(data, r - 1, i + 1, progress + [data[i]], result)

result = combinationListRecursive(range(len(d)), k)
best = 0
for i in result:
    tarray = []
    dtmp = 0
    for k in range(k):
        dtmp += d[i[k]]
        tarray.append(t[i[k]])
    tlist = list(set(tarray))
    ttmp = len(tlist) ** 2
    if best < dtmp + ttmp:
        best = dtmp + ttmp
print(best)

from heapq import heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()
for i in range(len(TD)):
    flg = 0
    if TD[i][0] not in muptiple:
        muptiple.add(TD[i][0])
        flg = 1
    TD[i][0] = flg
TD.reverse()
ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = TD.pop()
    vari += v[0]
    point += v[1]
    if not v[0]:
        heappush(ret, v[1])

total = vari ** 2 + point

while len(TD) > 0 and len(ret) > 0:
    v = TD.pop()
    if v[0] == 0:
        continue
    pre_t = heappop(ret)
    point -= pre_t - v[1]
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)

Verdicts: Runtime Error

Runtime Error

Runtime Error

Accepted