Problem:


Score : 400 points

Problem Statement

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.

Constraints

1 \leq K \leq N  \leq 10^5

1 \leq t_i \leq N

1 \leq d_i \leq 10^9

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

Output

Print the maximum satisfaction that you can obtain.

Sample Input 1

5 3
1 9
1 7
2 6
2 5
3 1

Sample Output 1

26

If you eat Sushi 1,2 and 3:

The base total deliciousness is 9+7+6=22.

The variety bonus is 2*2=4.

Thus, your satisfaction will be 26, which is optimal.

Sample Input 2

7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5

Sample Output 2

25

It is optimal to eat Sushi 1,2,3 and 4.

Sample Input 3

6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000

Sample Output 3

4900000016

Note that the output may not fit into a 32-bit integer type.



Sample input:
5 3
1 9
1 7
2 6
2 5
3 1


Sample output:
26

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
sushi = []
for i in range(n):
    tmp = list(map(int, input().split()))
    sushi.append(tmp)
eaten = set() # 食べた種類を記憶(集合型)
chohuku = [] # 同じ種類で二回以上食べた寿司のスコアを記録
scores = [] # Sのスコアのリスト

sushi.sort(key=lambda x:x[1])
sushi.reverse() # sushiをおいしさの降順にソート

tmp = 0 # Sのおいしさの合計を保持
for i in range(k): # sushiを美味しい順にk個食べる
    if sushi[i][0] in eaten: # 同じ種類の寿司を食べていた場合
        chohuku.append(sushi[i][1])
    else:
        eaten.add(sushi[i][0]) # 食べた寿司の種類を記憶
    tmp += sushi[i][1] # 食べた寿司のおいしさをtmpに追加
chohuku.sort()
scores.append(tmp+len(eaten)**2) # 今のSのスコアを記憶

for i in range(k,n):
    if not sushi[i][0] in eaten:
        tmp -= chohuku.pop()
        eaten.add(sushi[i][0])
        tmp += sushi[i][1]
        scores.append(tmp+len(eaten)**2)
        
print(max(scores))

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
sushi = []
for i in range(n):
    tmp = list(map(int, input().split()))
    sushi.append(tmp)
eaten = set() # 食べた種類を記憶(集合型)
chohuku = [] # 同じ種類で二回以上食べた寿司のスコアを記録
scores = [] # Sのスコアのリスト

sushi.sort(key=lambda x:x[1])
sushi.reverse() # sushiをおいしさの降順にソート

tmp = 0 # Sのおいしさの合計を保持
for i in range(k): # sushiを美味しい順にk個食べる
    if sushi[i][0] in eaten: # 同じ種類の寿司を食べていた場合
        chohuku.append(sushi[i][1])
    else:
        eaten.add(sushi[i][0]) # 食べた寿司の種類を記憶
    tmp += sushi[i][1] # 食べた寿司のおいしさをtmpに追加
chohuku.sort()
scores.append(tmp+len(eaten)**2) # 今のSのスコアを記憶
print(chohuku)

for i in range(k,n):
    if not sushi[i][0] in eaten:
        tmp -= chohuku.pop()
        eaten.add(sushi[i][0])
        tmp += sushi[i][1]
        scores.append(tmp+len(eaten)**2)
        
print(max(scores))

Submission status: Runtime Error

Solution code:
from collections import defaultdict

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

sushi = sorted(sushi, key = lambda x: x[1], reverse = True)
types = set()
max_vals = [0]
other_vals = [0]

for typ, val in sushi:
    if typ not in types:
        max_vals.append(max_vals[-1] + val)
        types.add(typ)
    else:
        other_vals.append(other_vals[-1] + val)

max_score = 0
for k in range(1, K + 1):
    if K - k < len(other_vals):
        score = max_vals[k] + other_vals[K - k] + k ** 2
        max_score = max(max_score, score)

print(max_score)

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
typ = set()
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.append(item[0])
        item[0] = 1
        lstA.append(item[1])
ans = 0
neta = 0
for i in range(K):
    ans += sushi[i][1]
    if sushi[i][0] == 1:
        neta += 1
ans += neta ** 2
C = sum(lstA[:neta])
D = sum(lstB[:K - neta])
for i in range(neta + 1, min(K, len(lstA)) + 1):
    C += lstA[i - 1]
    D -= lstB[K - i]
    tmp = C + D + i ** 2
    ans = max(ans, tmp)
print(ans)

Submission status: Runtime Error

Solution code:
# -*- 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()
# greedy for delicious point
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

# reduce dericious, gain variety
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)

    if len(pop_cands) == 0:
        break

print(mp_max)

Submission status: Runtime Error

Solution code:
# -*- 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()
# greedy for delicious point
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

# reduce dericious, gain variety
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)

Submission status: Runtime Error

Solution code:
class sushi:
  def __init__(self, t, d):
    self.t = t
    self.d = d

list = []
N, K = map(int, input().split())
for i in range(N):
  t, d = map(int, input().split())
  list.append(sushi(t, d))
  
sorted_list = sorted(list, key=lambda sushi:sushi.d, reverse=True)

t_list = []
rev_list = []
d_sum = 0
for i in range(K):
  if t_list.count(list[i].t) == 0:
    t_list.append(list[i].t)
  else:
    rev_list.append(i)
  d_sum += list[i].d

tmp = pow(len(t_list), 2) + d_sum

flag = 0
while len(rev_list) != 0:
  pointer = rev_list.pop()
  for i in range(K, N, 1):
    if t_list.count(list[i].t) == 0
      l = len(t_list)
      key = pow(l + 1, 2) + list[i].d - pow(l, 2) - list[pointer].d
      if key > 0:
        tmp += key
      else:
        flag = 1
      break
  if flag:
    break

print(tmp)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
from sys import stdin
from collections import Counter
import numpy as np


s_in = lambda: stdin.readline()[:-1]  # s = s_in()
d_in = lambda: int(stdin.readline())  # N = d_in()
ds_in = lambda: list(map(int, stdin.readline().split()))  # List = ds_in()


N, K = ds_in()
t_list = []
d_list = []
for _ in range(N):
    t, d = ds_in()
    t_list.append(t)
    d_list.append(d)
np_t = np.array(t_list)
np_d = np.array(d_list)

sort_idx = np.argsort(np_d)[::-1][:K]
top_sushi = np_t[sort_idx]
top_score = np_d[sort_idx]
kind = len(set(top_sushi))
point = top_score.sum() + kind**2
point_list = [point]
sushi_counter = Counter(top_sushi)

sushi_set = {k: [] for k in set(t_list)}
for i, j in zip(t_list, d_list):
    sushi_set[i].append(j)
max_sushi = {k: max(sushi_set[k]) for k in set(t_list)}

for s in set(top_sushi):
    max_sushi.pop(s)
max_values = list(max_sushi.values())
max_values.sort(reverse=True)

for sushi, score in zip(top_sushi[::-1], top_score[::-1]):
    if sushi_counter[sushi] > 1:
        sushi_counter[sushi] -= 1
        kind += 1
        point = point - score + max_values.pop(0) + 2*kind - 1
        point_list.append(point)
        if len(max_values) == 0:
            break

print(max(point_list))

Submission status: Runtime Error

Solution code:
#ABC116D-VariousSushi.py

from sys import stdin
from numpy import sum

input = stdin.readline

N, K = [int(n) for n in input().split()]
INPUT = [[None for i in range(2)] for j in range(N)]
List1, List2 = [], []
kind1 = []
flag1 = 0

for i in range(N):
    INPUT[i][0], INPUT[i][1] = [int(m) for m in input().split()]

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

for i in range(N):
    if INPUT[i][0] not in kind1:
            List1.append([INPUT[i][1], 1])
            kind1.append(INPUT[i][0])
    else: List1.append([INPUT[i][1], 0])

List2 = List1[K:]
if List2 != []:
    List2.sort(key =lambda x : (x[1],x[0]))
    del List1[K:]
    lenList2 = len(List2)
else: lenList2 = 0

List1 = sorted(List1, key = lambda x : (x[1], x[0]), reverse=True)

sum = sum(List1, axis = 0)

Score = sum[0] + sum[1]**2
Test = Score
kindnumber = sum[1]

for i in range(lenList2):
    if List2[-1][1] == 0 or List1[-1] == 1:
        break
    Test = Test - List1.pop()[0] + List2.pop()[0] + 2 * kindnumber + 1
    kindnumber = kindnumber + 1
    if Test >= Score:
        Score = Test
    i = i - 1

print(int(Score))

Submission status: Runtime Error

Solution code:
#ABC116D-VariousSushi.py

from sys import stdin
from numpy import sum

input = stdin.readline

N, K = [int(n) for n in input().split()]
INPUT = [[None for i in range(2)] for j in range(N)]
List1, List2 = [], []
kind1 = []
flag1 = 0

for i in range(N):
    INPUT[i][0], INPUT[i][1] = [int(m) for m in input().split()]

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

for i in range(N):
    if INPUT[i][0] not in kind1:
            List1.append([INPUT[i][1], 1])
            kind1.append(INPUT[i][0])
    else: List1.append([INPUT[i][1], 0])

List2 = List1[K:]
List2.sort(key =lambda x : (x[1],x[0]))
del List1[K:]
List1 = sorted(List1, key = lambda x : (x[1], x[0]), reverse=True)

sum = sum(List1, axis = 0)

Score = sum[0] + sum[1]**2
Test = Score
lenList2 = len(List2)
kindnumber = sum[1]

for i in range(lenList2):
    if List2[-1][1] == 0 or List1[-1] == 1:
        break
    Test = Test - List1.pop()[0] + List2.pop()[0] + 2 * kindnumber + 1
    kindnumber = kindnumber + 1
    if Test >= Score:
        Score = Test
    i = i - 1

print(int(Score))

Submission status: Runtime Error

Solution code:
#ABC116D-VariousSushi.py

import sys
import numpy as np

input = sys.stdin.readline

N, K = [int(n) for n in input().split()]
INPUT = [[None for i in range(2)] for j in range(N)]
List1, List2 = [], []
kind1, kind2 = [], []

for i in range(N):
    INPUT[i][0], INPUT[i][1] = [int(m) for m in input().split()]

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

for i in range(N):
    if INPUT[i][0] not in kind1:
        if len(List1) < K:
            List1.append(INPUT[i][1])
            kind1.append(INPUT[i][0])
    elif INPUT[i][0] not in kind2:
        List2.append(INPUT[i][1])
        kind2.append(INPUT[i][0])

List2.sort()

kindnumber = len(kind1)
k = len(List2)-(K-kindnumber)

Score = np.sum(List1) + np.sum(List2[k:]) + kindnumber ** 2
del List2[k:]
Test = Score

for i in range(kindnumber):
    Test = Test - List1.pop() + List2.pop() - 2 * kindnumber + 1
    kindnumber = kindnumber - 1
    if Test >= Score:
        Score = Test
    if List2 == []:
        break

print(int(Score))

Submission status: Runtime Error

Solution code:
#ABC116D-VariousSushi.py

import sys
import numpy as np
from operator import itemgetter

input = sys.stdin.readline

INPUT1 = [int(n) for n in input().split()]

N, K = INPUT1[0], INPUT1[1]
List1, List2 = [], []
Score= 0
kind = []

for i in range(N):
    INPUT2 = [int(m) for m in input().split()]
    if INPUT2[0] not in kind and len(List1) <= K:
        List1.append(INPUT2[1])
        kind.append(INPUT2[0])
    elif INPUT2[0] in kind:
        List2.append(INPUT2[1])

List1.sort(reverse=True)
List2.sort()

kindnumber = len(List1)
k = len(List2)-(K-kindnumber)

Score = np.sum(List1) + np.sum(List2[k:]) + kindnumber ** 2
del List2[k:]
Test = Score

for i in range(kindnumber):
    Test = Test - List1.pop() + List2.pop() - 2 * kindnumber + 1
    kindnumber = kindnumber - 1
    if Test >= Score:
        Score = Test
    if List2 == []:
        break

print(int(Score))

Submission status: Runtime Error

Solution code:
import itertools

a = input().split()
osu = []
for i in range(int(a[0])):
    osu.append(input().split())

im = []
for i in range(int(a[0])):
    im.append(i)
ans_map_all = list(itertools.combinations(im, int(a[1])))

ans_lis = []
for i in ans_map_all:
    kind = []
    pre = 0
    for j in i:
        kind.append(int(osu[j][0]))
        pre += int(osu[j][1])
    ans_lis.append(pre + len(list(set(kind))*len(list(set(kind)))))

ans_lis.sort()
ans_lis.reverse()
print(ans_lis[0])

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1], reverse=True)

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = set()
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types = types.union({c[0]})

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for i, l in enumerate(leave):
    if i > types_max:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i < types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      leave = leave[:j]
      curr += [leave.pop(0)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-


if __name__=='__main__':
  N, K = map(int, input().split())
  sushis = sorted([list(map(int, input().split())) for _ in range(N)], key=lambda x:x[1])[::-1]

  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types = set([c[0] for c in curr])
  di = {}

  types_max = len(set([c[0] for c in sushis]))

  for i in range(len(types), K+1):
    di[str(i)] = base(curr) + i**2

    # next
    if i <= types_max:
      tdi = {}
      remove_number = -1
      for j, t in enumerate(curr[::-1]):
        if not str(t[0]) in tdi.keys():
          tdi[str(t[0])] = -j-1
        else:
          remove_number = tdi[str(t[0])]
          break
      curr.pop(int(remove_number))

      for j, l in enumerate(leave):
        if not l[0] in types:
          types.union({l[0]})
          break
      curr += [leave.pop(j)]
    else:
      break
  print(sorted(di.items(), key=lambda x:x[1])[-1][1])

Submission status: Runtime Error

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    TD = []
    for i in range(N):
        TD.append(LI())
    #  print(T, '\n', D)
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    results = []
    for i in range(1, K+1):
        if i == K:
            # 全タイプの最大値を取る
            base_point = 0
            for t in [set(i[0] for i in TD)]:
                base_point += max([i for i in TD if i[0] == t])

            results.append(base_point + i * i)


        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 最小種類数のときの最大値、最大種類数のときの最大値を計算して比較する
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(set([i[0] for i in TD]), i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            base_point = 0
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            required_neta_count = K
            # 各タイプの最大値をまず選び、current_optionsから排除
            for t in combination:
                max_value = max([option[1] for option in current_options if option[0] == t])
                base_point += max_value

                # max_valueは選択肢から取り除いておく
                current_options = [i for i in current_options if not (i[0] == t and i[1] == max_value)]
                #  print('選んだ寿司を選択肢から取り除いた', current_options)

            #  print('あと何個ネタを取るべきか? ', required_neta_count, '個')
            # あとは残りの選択肢から大きいものを選ぶ
            #  print('current_options, requried_neta_count', current_options, required_neta_count)
            results.append(sum([i[1] for i in current_options[:K-i]]) + base_point + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()

Submission status: Runtime Error

Solution code:
import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
from operator import itemgetter


sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)

def main():
    # 種類最小から最大までそれぞれ計算する
    N, K = LI()
    T = []
    D = []
    for i in range(N):
        t, d = LI()
        T.append(t)
        D.append(d)
    #  print(T, '\n', D)
    TD = [[t, d] for t, d in zip(T, D)]
    TD = sorted(TD, key=itemgetter(1), reverse=True)
    #  print('TD', TD)

    all_types = set(T)
    results = []
    for i in range(1, K+1):
        #  print('\n---------\n',i,'回目')
        # iを種類数とみなす
        # 最小種類数のときの最大値、最大種類数のときの最大値を計算して比較する
        # 何を選んだかは関係ない
        
        # 種類の全組み合わせを求める
        #  print(i)
        combinations = list(itertools.combinations(all_types, i))
        #  print('combi', combinations)
        type_bonus = i * i
        for combination in combinations:
            #  print('\n')
            base_point = 0
            # 各組み合わせでの最大値を求める
            current_options = [i for i in TD if i[0] in combination]
            #  print('current_options: ', current_options)
            # そもそも数がKに満たない場合は作るのが不可能なため次のループへ
            if len(current_options) < K:
                continue

            required_neta_count = K
            # 各タイプの最大値をまず選び、current_optionsから排除
            for t in combination:
                base_point += max([option[1] for option in current_options if option[0] == t])

                # max_valueは選択肢から取り除いておく
                current_options = [i for i in current_options if not (i[0] == t and i[1] == max_value)]
                #  print('選んだ寿司を選択肢から取り除いた', current_options)

            #  print('あと何個ネタを取るべきか? ', required_neta_count, '個')
            # あとは残りの選択肢から大きいものを選ぶ
            #  print('current_options, requried_neta_count', current_options, required_neta_count)
            results.append(sum([i[1] for i in current_options[:K-i]]) + base_point + type_bonus)

    #  print('\n======================\nresults', results)
    print(max(results))


main()

Submission status: Runtime Error

Solution code:
#include <bits/stdc++.h>
using ll = long long;
#define int ll
using namespace std;
int N, K;
priority_queue<pair<int, int>> q;
signed main()
{
  cin >> N >> K;
  for (int i = 0; i < N; i++){
    int t, d;
    cin >> t >> d;
    q.push(make_pair(d, t));
  }
  vector<int> c(N + 1, 0);
  ll k = 0, res = 0;
  priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> mi;
  for (int i = 0; i < K; i++){
    auto x = q.top();
    q.pop();
    res += x.first;
    if (c[x.second] == 0){
      k++;
    }
    mi.push(x);
    c[x.second]++;
  }
  ll ma = res + k * k;
  for (int i = k + 1; i < K + 1; i++){
    while (c[mi.top().second] == 1) {
      mi.pop();
    }
    res -= mi.top().first;
    c[mi.top().second]--;
    mi.pop();
    while(!q.empty() && c[q.top().second]){
      q.pop();
    }
    if (q.empty()) {
      break;
    }
    res += q.top().first;
    c[q.top().second]++;
    q.pop();
    ma = max(ma, res + i * i);
  }
  cout << ma << endl;
}

Submission status: Runtime Error

Solution code:
import sys

def point(sushi):
  d = 0
  t = []
  for s in sushi:
    d += s[1]
    t.append(s[0])
  return d + len(set(t))**2


def rm_ate(sushi, rest):
  t = set([s[0] for s in sushi])
  return [r for r in rest if not r[0] in t]


def rm_dup(sushi):
  sushi = sorted(sushi, key=lambda x:-x[1])
  d = []
  r = []
  for s in sushi:
    if s[0] in d:
      r.append(s)
    else:
      d.append(s[0])
  sushi.remove(r[-1])
  return sushi


def main():
  input = sys.stdin.readline
  n, k = map(int, input().split())
  sushi = []
  for _ in range(n):
    t, d = map(int, input().split())
    sushi.append([t,d])
  sushi = sorted(sushi, key=lambda x:x[1])
  temp = sushi[-k:]
  rest = rm_ate(temp, sushi[:-k])
  ans = point(temp)
  while rest:
    temp = rm_dup(temp)
    temp += [rest.pop()]
    res = point(temp)
    ans = max(ans, res)
    rest = rm_ate(temp, rest)
  print(ans)


if __name__ == '__main__':
  main()

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
# Luis Yoichi Morales January 2019

from sys import stdin
import numpy as np

def compute_satisfaction (K, h):
    d_sum=0
    types=[]
    for i in range (K):
        d_sum+=h[i][1]
        types.append( h[i][0] )

    np_types = np.array( types )
    taste_diff_types = len( ( np.unique ( np_types ) ) )
    return long(d_sum + taste_diff_types * taste_diff_types)



N,K = map(int, stdin.readline().split())
h=[]
for i in range (N):
    h.append( map(int, stdin.readline().split()) )

#Order descening order for taste
h.sort(key=lambda x: -x[1])

#Get larger delicious sum
taste_satisfaction = compute_satisfaction (K, h) 


#Get larger score for varied menu
np_types = np.array( h )
taste_diff_types = len( ( np.unique ( np_types ) ) )

if taste_diff_types > K:    
    diff=[]
    diff.append(h[0])
    for i in range (1, N):
        if h[i][0] != h[i-1][0]:
            diff.append(h[i])
            if len(diff) == K:
                break
    
    different_menu_satisfaction =  compute_satisfaction (K, diff)
else:
    different_menu_satisfaction =  0

if taste_satisfaction > different_menu_satisfaction:
    print taste_satisfaction
else:
    print different_menu_satisfaction

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
# Luis Yoichi Morales January 2019

from sys import stdin
import numpy as np

def compute_satisfaction (K, h):
    d_sum=0
    types=[]
    for i in range (K):
        d_sum+=h[i][1]
        types.append( h[i][0] )

    np_types = np.array( types )
    taste_diff_types = len( ( np.unique ( np_types ) ) )
    return d_sum + taste_diff_types * taste_diff_types



N,K = map(int, stdin.readline().split())
h=[]
for i in range (N):
    h.append( map(int, stdin.readline().split()) )

#Order descening order for taste
h.sort(key=lambda x: -x[1])

#Get larger delicious sum
taste_satisfaction = compute_satisfaction (K, h) 


#Get larger score for varied menu
np_types = np.array( h )
taste_diff_types = len( ( np.unique ( np_types ) ) )

if taste_diff_types > K:    
    diff=[]
    diff.append(h[0])
    for i in range (1, N):
        if h[i][0] != h[i-1][0]:
            diff.append(h[i])
            if len(diff) == K:
                break
    
    different_menu_satisfaction =  compute_satisfaction (K, diff)
else:
    different_menu_satisfaction =  0

if taste_satisfaction > different_menu_satisfaction:
    print taste_satisfaction
else:
    print different_menu_satisfaction

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
# Luis Yoichi Morales January 2019

from sys import stdin
import numpy as np

def compute_satisfaction (K, h):
    d_sum=0
    types=[]
    for i in range (K):
        d_sum+=h[i][1]
        types.append( h[i][0] )

    np_types = np.array( types )
    taste_diff_types = len( ( np.unique ( np_types ) ) )
    return d_sum + taste_diff_types * taste_diff_types



N,K = map(int, stdin.readline().split())
h=[]
for i in range (N):
    h.append( map(int, stdin.readline().split()) )

#Order descening order for taste
h.sort(key=lambda x: -x[1])




#Get larger delicious sum
np_types = np.array( h )
taste_diff_types = len( ( np.unique ( np_types ) ) )

if taste_diff_types > K:

    d_sum=0
    types=[]
    taste_satisfaction = compute_satisfaction (K, h) 

    #Get larger score for varied menu
    diff=[]
    diff.append(h[0])
    for i in range (1, N):
        if h[i][0] != h[i-1][0]:
            diff.append(h[i])
            if len(diff) == K:
                break
    
    different_menu_satisfaction =  compute_satisfaction (K, diff)
else:
    different_menu_satisfaction =  0

if taste_satisfaction > different_menu_satisfaction:
    print taste_satisfaction
else:
    print different_menu_satisfaction

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
# Luis Yoichi Morales January 2019

from sys import stdin
import numpy as np

def compute_satisfaction (K, h):
    d_sum=0
    types=[]
    for i in range (K):
        d_sum+=h[i][1]
        types.append( h[i][0] )

    np_types = np.array( types )
    taste_diff_types = len( ( np.unique ( np_types ) ) )
    return d_sum + taste_diff_types * taste_diff_types



N,K = map(int, stdin.readline().split())
h=[]
for i in range (N):
    h.append( map(int, stdin.readline().split()) )

#Order descening order for taste
h.sort(key=lambda x: -x[1])



#Get larger delicious sum
d_sum=0
types=[]
taste_satisfaction = compute_satisfaction (K, h) 

#Get larger score for varied menu
diff=[]
diff.append(h[0])
for i in range (1, N):
    if h[i][0] != h[i-1][0]:
        diff.append(h[i])
    if len(diff) == K:
        break

different_menu_satisfaction =  compute_satisfaction (K, diff)

if taste_satisfaction > different_menu_satisfaction:
    print taste_satisfaction
else:
    print different_menu_satisfaction

Submission status: Runtime Error

Solution code:
import sys

N,K=map(int, sys.stdin.readline().split()) 
T=[ map(int, sys.stdin.readline().split()) for i in range(N) ]
T=sorted(T, key=lambda x:(x[0],-x[1]))

KI=[]	#KIND
A=[]	#Nokori

x=None
for t,d in T:
	if t != x:
		KI.append(d)
		x=t
	else:
		A.append(d)

KI.sort(reverse=True)
A.sort(reverse=True)

KI_S=[0]		#KI Ruiseki sum
A_S=[0]		#A Ruiseki sum

for x in KI: KI_S.append( KI_S[-1]+x )
for x in A:  A_S.append( A_S[-1]+x )


l=len(A_S)

ans=0
for k in range(1,K+1):
	if K-k<l:
		ans=max (ans , KI_S[k]+k**2+A_S[K-k] )

print ans

Submission status: Runtime Error

Solution code:
from collections import Counter as ct
n, k = map(int, input().split())

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

kind = len(ct(list(map(lambda x: x[0], td))))
std = sorted(td, key=lambda x: x[1], reverse=True)
dup = []
eat_k = set()
ans = 0
for i in range(k):
    ans += std[i][1]
    if std[i][0] in eat_k:
        dup.append(std[i])
    eat_k.add(std[i][0])
ans += len(eat_k)** 2
tmpa = ans
for i in range(k, n):
    if std[i][0] not in eat_k:
        t = dup.pop()
        tmpa -= t[1]
        tmpa += std[i][1]
        tmpa -= len(eat_k)** 2
        eat_k.add(std[i][0])
        tmpa += len(eat_k)** 2
        ans = max(tmpa, ans)
print(ans)

Submission status: Runtime Error

Solution code:
from collections import Counter as ct
n, k = map(int, input().split())

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

kind = len(ct(list(map(lambda x: x[0], td))))
std = sorted(td, key=lambda x: x[1], reverse=True)
dup = []
eat_k = set()
ans = 0
for i in range(k):
    ans += std[i][1]
    if std[i][0] in eat_k:
        dup.append(std[i])
    eat_k.add(std[i][0])
ans += len(eat_k)** 2
tmpa = ans
for i in range(n - k, n):
    if std[i][0] not in eat_k:
        t = dup.pop()
        tmpa -= t[1]
        tmpa += std[i][1]
        tmpa -= len(eat_k)** 2
        eat_k.add(std[i][0])
        tmpa += len(eat_k)** 2
        ans = max(tmpa, ans)

print(ans)