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

inf = float('inf')

n, k = map(int, input().split())

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

for key in dd:
    dd[key].sort(reverse=True)

h1 = []
h2 = []
for key in dd:
    for i, d in enumerate(dd[key]):
        if i == 0:
            h1.append(d)
        else:
            h2.append(d)

if h1:
    h1.sort(reverse=True)
if h2:
    h2.sort(reverse=True)

c1 = [0]
for hh in h1:
    c1.append(c1[-1] + hh)
if len(c1) < k:
    c1.extend([-inf] * (k - len(c1) + 1))

c2 = [0]
for hh in h2:
    c2.append(c2[-1] + hh)
if len(c2) < k:
    c2.extend([-inf] * (k - len(c2) + 1))

# print(h1, c1)
# print(h2, c2)

ans = -inf
for kind in range(1, k + 1):
    # print(kind, c1[kind], c2[k - kind])
    ans = max(ans, c1[kind] + c2[k - kind] + kind * kind)
print(ans)

Submission status: Runtime Error

Solution code:
from collections import defaultdict

inf = float('inf')

n, k = map(int, input().split())

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

for key in dd:
    dd[key].sort(reverse=True)

h1 = []
h2 = []
for key in dd:
    for i, d in enumerate(dd[key]):
        if i == 0:
            h1.append(d)
        else:
            h2.append(d)

h1.sort(reverse=True)
h2.sort(reverse=True)

c1 = [0]
for hh in h1:
    c1.append(c1[-1] + hh)
if len(c1) < k:
    c1.extend([-inf] * (k - len(c1) + 1))

c2 = [0]
for hh in h2:
    c2.append(c2[-1] + hh)
if len(c2) < k:
    c2.extend([-inf] * (k - len(c2) + 1))

# print(h1, c1)
# print(h2, c2)

ans = -inf
for kind in range(1, k + 1):
    # print(kind, c1[kind], c2[k - kind])
    ans = max(ans, c1[kind] + c2[k - kind] + kind * kind)
print(ans)

Submission status: Runtime Error

Solution code:
import heapq

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

d_sushi = []  # 捨てる候補の寿司
heapq.heapify(d_sushi)  # リストaのheap化

fp = 0
tp = set()

for i in range(K):
    fp += sushi[i][1]
    if sushi[i][0] in tp:
        heapq.heappush(d_sushi, sushi[i][1])
    else:
        tp.add(sushi[i][0])
ans = fp+len(tp)**2

j = K
while d_sushi and len(tp) <= N:
    while sushi[j][0] in tp:
        j += 1
        if j == N:
            print(ans)
            exit()
    fp = fp-heapq.heappop(d_sushi)+sushi[j][1]
    tp.add(sushi[j][0])
    ans = max(ans, fp+len(tp)**2)

print(ans)

Submission status: Runtime Error

Solution code:
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)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10**7)
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 II(): return int(sys.stdin.readline())
def SI(): return sys.stdin.readline().strip()
INF = 10 ** 18
MOD = 10 ** 9 + 7

def score(td):
  t_list = [v[0] for v in td]
  d_list = [v[1] for v in td]
  ut = len(set(t_list))
  score_ut = ut**2
  score_d = sum(d_list)
  return score_ut + score_d
  
def solve():
  N, K = LI()
  td = []
  for i in range(N):
    t, d = LI()
    td.append([t, d])

  from operator import itemgetter
  td.sort(key=itemgetter(1), reverse=True)

  # おいしさtopK
  cand = td[:K]
  
  # cand のネタ保持数dict
  import collections
  c = collections.Counter([v[0] for v in cand])
  m = len(c.keys())
  
  # 持ってないネタのtop1リスト
  rest_top1 = {}
  for v in td:
    if v[0] not in c:
      if v[0] not in rest_top1:
      	rest_top1[v[0]] = v[1]
      else:
        if rest_top1[v[0]] < v[1]:
          rest_top1[v[0]] = v[1]
          
  # スコア降順
  rest_o = sorted([v for k, v in rest_top1.items()], reverse=True)
  
  current_score = score(cand)
  scores = [current_score]
  # おいしく無い順に持ってないもののtop1と入れ替え
  nget = 0
  for i in range(1, K):
    target = cand[-i]
    if c[target[0]] == 1:
      continue # 一つしか持ってないなら入れ替えない
    
    # スコアの差分を作る
    # netaボーナス
    dscore_neta = (m+nget+1)**2 - (m+nget)**2 
    # おいしさボーナス
    dscore_d = -target[1] + rest_o[nget]
    current_score += dscore_neta
    current_score += dscore_d
    
    scores.append(current_score)
    # 状態更新
    nget += 1
    c[target[0]] -= 1
    
    # 全部取ったらおわり
    if nget == len(rest_o):
      break
  print(max(scores))
solve()

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10**7)
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 II(): return int(sys.stdin.readline())
def SI(): return sys.stdin.readline().strip()
INF = 10 ** 18
MOD = 10 ** 9 + 7

def score(td):
  t_list = [v[0] for v in td]
  d_list = [v[1] for v in td]
  ut = len(set(t_list))
  score_ut = ut**2
  score_d = sum(d_list)
  return score_ut + score_d
  
def solve():
  N, K = LI()
  td = []
  for i in range(N):
    t, d = LI()
    td.append([t, d])

  from operator import itemgetter
  td.sort(key=itemgetter(1), reverse=True)

  # おいしさtopK
  cand = td[:K]
  
  # cand のネタ保持数dict
  import collections
  c = collections.Counter([v[0] for v in cand])
  m = len(c.keys())
  
  # 持ってないネタのtop1リスト
  rest_top1 = {}
  for v in td:
    if v[0] not in c:
      if v[0] not in rest_top1:
      	rest_top1[v[0]] = v[1]
      else:
        if rest_top1[v[0]] < v[1]:
          rest_top1[v[0]] = v[1]
          
  # スコア降順
  rest_o = sorted([v for k, v in rest_top1.items()], reverse=True)
  
  current_score = score(cand)
  scores = [current_score]
  # おいしく無い順に持ってないもののtop1と入れ替え
  nget = 0
  for i in range(1, K):
    target = cand[-i]
    if c[target[0]] == 1:
      continue # 一つしか持ってないなら入れ替えない
    
    # スコアの差分を作る
    # netaボーナス
    dscore_neta = (m+nget+1)**2 - (m+nget)**2 
    # おいしさボーナス
    dscore_d = -target[1] + rest_o[nget]
    current_score += dscore_neta
    current_score += dscore_d
    
    scores.append(current_score)
    # 状態更新
    nget += 1
    c[target[0]] -= 1
  print(max(scores))
solve()

Submission status: Runtime Error

Solution code:
def main():
    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)

    first = set()
    second = []
    point = 0
    for s in sushi[:k]:
        if s[0] not in first:
            first.add(s[0])
        else:
            second.append(s)
        point += s[1]
    point += len(first) ** 2
    ans = point

    for s in sushi[k:]:
        if s[0] not in first:
            point = point + s[1] - second.pop()[1] + \
                (len(first) + 1) ** 2 - len(first) ** 2
            first.add(s[0])
            ans = max(ans, point)
    print(ans)


main()

Submission status: Runtime Error

Solution code:
from collections import defaultdict
from operator import itemgetter
import sys

input = sys.stdin.readline

n, k = map(int, input().split())
e = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    e[t].append(d)

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, key=itemgetter(0), reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)

Submission status: Runtime Error

Solution code:
import heapq

N, K = map(int, input().split())
sushi = []
for i in range(N):
    t, d = map(int, input().split())
    heapq.heappush(sushi, (-d, t))

total = 0
types = set()
remain = []
for i in range(K):
    d, t = heapq.heappop(sushi)
    total -= d
    if t not in types:
        types.append(t)
    else:
        # 余りの寿司はremainにappend
        heapq.heappush(remain, -d)

types_len = len(types)
ans = total + types_len ** 2

while sushi and remain:
    d, t = heapq.heappop(sushi)
    if t not in types:
        types.append(t)
        remain_d = heapq.heappop(remain)
        total = total - d - remain_d
        types_len += 1
        ans = max(ans, total + types_len ** 2)

print(ans)

Submission status: Runtime Error

Solution code:
from operator import itemgetter
n,k = map(int,input().split())
li=[]
s=[]
t=[]
g=0
ans=[]
for _ in range(n):
    a,b=map(int,input().split())
    li.append([a,b])
sorted(li,key=itemgetter(1),reverse=True)
for i in range(k):
    g+=li[i][1]
    if li[i][0] not in s:
        s.append(li[i][0])
    else:
        t.append(li[i][1])
g+=len(s)**2
ans.append(g)
for i in range(k,n):
    if t==0:
        break
    if li[i][0] not in s:
        g=g-t[-1]+li[i][1]+2*len(s)+1
        ans.append(g)
        s.append(li[i][0])
print(max(ans))

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)
#'''


def caluculate_score():
	return num_kind * num_kind + taste
	
	
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
table = [[] for _ in range(N)]
eat_kind = [0] * N
num_kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	table[i] = (d[i], t[i] - 1)
	
table.sort()

eat = []
taste = 0

for _ in range(K):
	tmp = table.pop()
	taste += tmp[0]
	eat_kind[tmp[1]] += 1
	if eat_kind[tmp[1]] == 1:
		num_kind += 1
	eat.append(tmp)
	
ans = caluculate_score()
ind_eat = K - 1
ind_table = N - K - 1

while ind_eat >= 0 and ind_table >= 0:
	while eat_kind[eat[ind_eat][1]] <= 1:
		ind_eat -= 1
		
	while eat_kind[table[ind_table][1]] >= 1:
		ind_table -= 1
		
	if not(ind_eat >= 0 and ind_table >= 0):
		break
		
'''
	eat_kind[eat[ind_eat][1]] -= 1
	eat_kind[table[ind_table][1]] += 1
	num_kind += 1
	taste += table[ind_table][0] - eat[ind_eat][0]
	eat[ind_eat], table[ind_table] = table[ind_table], eat[ind_eat]
	score = caluculate_score()
	if score > ans:
		ans = score
		
	ind_eat -= 1
	ind_table -= 1
'''
print(ans)

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)
#'''


def caluculate_score():
	return num_kind * num_kind + taste
	
	
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
table = [[] for _ in range(N)]
eat_kind = [0] * N
num_kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	table[i] = (d[i], t[i] - 1)
	
table.sort()

eat = []
taste = 0

for _ in range(K):
	tmp = table.pop()
	taste += tmp[0]
	eat_kind[tmp[1]] += 1
	if eat_kind[tmp[1]] == 1:
		num_kind += 1
	eat.append(tmp)
	
ans = caluculate_score()
ind_eat = K - 1
ind_table = N - K - 1

while ind_eat >= 0 and ind_table >= 0:
	while eat_kind[eat[ind_eat][1]] <= 1:
		ind_eat -= 1
		
	while eat_kind[table[ind_table][1]] >= 1:
		ind_table -= 1
		
	if not(ind_eat >= 0 and ind_table >= 0):
		break
		
	eat_kind[eat[ind_eat][1]] -= 1
	eat_kind[table[ind_table][1]] += 1
	num_kind += 1
	taste += table[ind_table][0] - eat[ind_eat][0]
	eat[ind_eat], table[ind_table] = table[ind_table], eat[ind_eat]
	score = caluculate_score()
	if score > ans:
		ans = score
		
	ind_eat -= 1
	ind_table -= 1
	
print(ans)

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)
#'''


def caluculate_score():
	return num_kind * num_kind + taste
	
	
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
table = [[] for _ in range(N)]
eat_kind = [0] * N
num_kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	table[i] = (d[i], t[i] - 1)
	
table.sort()

eat = []
taste = 0

for _ in range(K):
	tmp = table.pop()
	taste += tmp[0]
	eat_kind[tmp[1]] += 1
	if eat_kind[tmp[1]] == 1:
		num_kind += 1
	eat.append(tmp)
	
ans = caluculate_score()
ind_eat = K - 1
ind_table = N - K - 1

while ind_eat >= 0 and ind_table >= 0:
	while eat_kind[eat[ind_eat][1]] <= 1:
		ind_eat -= 1
		
	while eat_kind[table[ind_table][1]] >= 1:
		ind_table -= 1
		
	eat_kind[eat[ind_eat][1]] -= 1
	eat_kind[table[ind_table][1]] += 1
	num_kind += 1
	taste += table[ind_table][0] - eat[ind_eat][0]
	eat[ind_eat], table[ind_table] = table[ind_table], eat[ind_eat]
	score = caluculate_score()
	if score > ans:
		ans = score
		
	ind_eat -= 1
	ind_table -= 1
	
print(ans)

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
sushi = [[] for i in range(N+1)]
for _ in range(N):
    t, d = map(int, input().split())
    sushi[t].append(-d)

sushi.sort()
for i in range(N):
    if len(sushi[i]) > 0:
        sushi = sushi[i:]
        break

kinds_min = 0
count = 0
for i in range(N):
    count += len(sushi[i])
    if count >= K:
        kinds_min = i+1
        break

kinds_max = min(K, len(sushi))


import heapq

sushi
ans = 0
queue = []
for i in range(kinds_min-1):
    ans -= sushi[i][0]
    queue += sushi[i][1:]

heapq.heapify(queue)
ansmax = 0
ans
sushi



for x in range(kinds_min-1, kinds_max):
    for i in sushi[x][1:]:
        heapq.heappush(queue, i)
    ans -= sushi[x][0]
    sushilist = []
    for i in range(K-x-1):
        sushilist.append(heapq.heappop(queue))
    ans -= sum(sushilist)
    ansmax = max(ansmax, ans+(x+1)**2)
    ans += sum(sushilist)
    for i in sushilist:
        heapq.heappush(queue, i)

print(ansmax)

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
sushi = [[] for i in range(N+1)]
for i in range(N):
    t, d = map(int, input().split())
    sushi[t].append(-d)

sushi.sort()
for i in range(N):
    if len(sushi[i]) > 0:
        sushi = sushi[i:]
        break

kinds_min = 0
count = 0
for i in range(N):
    count += len(sushi[i])
    if count >= K:
        kinds_min = i+1
        break

kinds_max = min(K, len(sushi))


import heapq

sushi
ans = 0
queue = []
for i in range(kinds_min-1):
    ans -= sushi[i][0]
    queue += sushi[i][1:]

heapq.heapify(queue)
ansmax = 0
ans
sushi



for x in range(kinds_min-1, kinds_max):
    for i in sushi[x][1:]:
        heapq.heappush(queue, i)
    ans -= sushi[x][0]
    sushilist = []
    for i in range(K-x-1):
        sushilist.append(heapq.heappop(queue))
    ans -= sum(sushilist)
    ansmax = max(ansmax, ans+(x+1)**2)
    ans += sum(sushilist)
    for i in sushilist:
        heapq.heappush(queue, i)
        
print(ansmax)

Submission status: Runtime Error

Solution code:
from collections import Counter

n,k = map(int, input().split())

sushis = [[0 for _ in range(2)] for _ in range(n)]
cnt_sushi_type = {}
for i in range(n):
  t,p = map(int, input().split())
  sushis[i][0] = -p #ポイント、大きいものから取り出したいのでマイナスにして入れる
  sushis[i][1] = t #種類

cnt_types = Counter() #寿司の種類覚えておく用のdict

sushis.sort(key=lambda x:x[0])

sushi_set = []
ans = 0
for p, t in sushis[:k]:
  #大きいものからk個選んで得点を計算
  ans += -p
  cnt_types[t] += 1
  sushi_set.append([p,t]) #集合を保存
  
ans += len(cnt_types) ** 2 #ネタ数＊ネタ数を加点

##ここまで貪欲選択

#x種類選んだ時の最大値
max_types = [0] * n
max_types[len(cnt_types)] = ans
for p, t in sushis[k:]: #これまで選んでないものを選択した場合についてみていく
  if t in cnt_types:
    continue #選んでない種類のネタを選びたいので、すでに選んでるネタであれば何もせず飛ばす
  else:
    #選んだ寿司集合にある一番得点の低い寿司を、現在見ている新ネタの寿司と入れ替える
    #実際に入れ替えるとめんどくさい（計算量的にappendとかしんどい）ので、ansとcnt_typesをいじることで調整する

    while sushi_set:
      ps, ts = sushi_set.pop() #小さい順で取り出す
      if cnt_types[ts] >= 2:
        #そのネタが2個以上あるなら一個減らして新セットを作る
        tmp =  max_types[len(cnt_types)] - (-ps) + (-p) + 2*len(cnt_types) + 1
        max_types[len(cnt_types)+1] = tmp
        cnt_types[t] += 1
        cnt_types[ts] -= 1
        break #次のまだ見ぬネタを探しに行く
        
#各種類選んだ最大値から最大のものを選択で終了
print(max(max_types))

Submission status: Runtime Error

Solution code:
import sys
import math 
ni = lambda: int(ns())
na = lambda: list(map(int, input().split()))
ns = lambda: input()
N,K = na()
S = [na() for i in range(N)]
S.sort(key=lambda x:x[1],reverse=True)
res = []
tmp=0
d=set([])
ds = []
for t,dd in S[:K]:
    tmp+=dd
    if t in d:
        ds.append(dd)
    else:
        d.add(t)
ans = tmp + len(d)**2
ax = ans

for t,dd in S[K:]:
    if len(d)==0:
        break
    if t not in d:
        ax +=2*len(d)+1
        ax += dd-ds.pop()
        d.add(t)
        ans = max(ax,ans)
print(ans)

Submission status: Runtime Error

Solution code:
import sys
import math 
ni = lambda: int(ns())
na = lambda: list(map(int, input().split()))
ns = lambda: input()
N,K = na()
S = [na() for i in range(N)]
S.sort(key=lambda x:x[1],reverse=True)
res = []
tmp=0
d=set([])
ds = []
for t,dd in S[:K]:
    tmp+=dd
    if t in d:
        ds.append(dd)
    else:
        d.add(t)
ans = tmp + len(d)**2
ax = ans

for t,dd in S[K:]:
    if t not in d:
        ax +=2*len(d)+1
        ax += dd-ds.pop()
        d.add(t)
        ans = max(ax,ans)
print(ans)

Submission status: Runtime Error

Solution code:
import sys
import math 
from heapq import *
ni = lambda: int(ns())
na = lambda: list(map(int, input().split()))
ns = lambda: input()
N,K = na()
S = [na() for i in range(N)]
S.sort(key=lambda x:x[1],reverse=True)
res = []
tmp=0
d={}
ans = [0]*N
for i in S[:K]:
    tmp+=i[1]
    if i[0] not in d:
        d[i[0]]=1
    else:
        d[i[0]]+=1
ans[len(d)]=tmp+len(d)**2
for i in S[:K]:
    if d[i[0]]>1:
        res.append([i[1], i[0]])
heapify(res)

for i in range(len(res)):
    dd,t = heappop(res)
    for j in S[K:]:
        if d[t]>1 and (j[0] not in d):
            d[t] -= 1
            d[j[0]] = 1
            heappush(res,[j[1],j[0]])
            ans[len(d)] = ans[len(d)-1]+(j[1]-dd)+(len(d)**2-(len(d)-1)**2)
            break
print(max(ans))

Submission status: Runtime Error

Solution code:
from collections import defaultdict

def getlist():
	return list(map(int, input().split()))
D = defaultdict(int)
V = []

N, K = getlist()
for i in range(N):
	t, d = getlist()
	V.append([d, t])

V = sorted(V)

ans = []
ex = []
k = 0
for i in range(N):
	x = V[-1 - i][1]
	if k < K:
		if D[x] == 0:
			ans.append(V[-1 - i])
			D[x] += 1
			k += 1
	else:
		ex.append(V[-1 - i])


if k < K:
	for i in range(K - k):
		ans.append(ex[i])
	eex = []
	for i in range(K - k, len(ex)):
		eex.append(ex[i])
else:
	eex = ex
ans = sorted(ans)
eex = sorted(eex)


A = 0
for i in range(K):
	A += ans[i][0]
A += k * k
n = min(k, N - K)
for i in range(n):
	if ans[i][0] + ((k - i) ** 2) - ((k - i - 1) ** 2) < eex[-1 - i][0]:
		A += eex[-1 - i][0] - ans[i][0] - ((k - i) ** 2) + ((k - i - 1) ** 2)
	else:
		break

print(A)

Submission status: Runtime Error

Solution code:
from collections import defaultdict

def getlist():
	return list(map(int, input().split()))
D = defaultdict(int)
V = []

N, K = getlist()
for i in range(N):
	t, d = getlist()
	V.append([d, t])

V = sorted(V)

ans = []
ex = []
k = 0
for i in range(N):
	x = V[-1 - i][1]
	if D[x] == 0:
		ans.append(V[-1 - i])
		D[x] += 1
		k += 1
	else:
		ex.append(V[-1 - i])
	if k == K:
		break

if k < K:
	for i in range(K - k):
		ans.append(ex[i])
	eex = []
	for i in range(K - k, len(ex)):
		eex.append(ex[i])
else:
	eex = ex
ans = sorted(ans)
eex = sorted(eex)


A = 0
for i in range(K):
	A += ans[i][0]
A += k * k
n = min(k, N - K)
for i in range(n):
	if ans[i][0] + ((k - i) ** 2) - ((k - i - 1) ** 2) < eex[-1 - i][0]:
		A += eex[-1 - i][0] - ans[i][0] - ((k - i) ** 2) + ((k - i - 1) ** 2)
	else:
		break

print(A)

Submission status: Runtime Error

Solution code:
import sys
import heapq
input = sys.stdin.readline

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

sushi.sort(reverse=True)
ans = 0
num_d = 0
num_t = dict()
for i in range(k):
    d, t = sushi[i]
    num_d += d
    if not t in num_t:
        num_t[t] = 1
    else:
        num_t[t] += 1
ans = max(ans, num_d + len(num_t) ** 2)
s = sushi[:k]
heapq.heapify(s)
ns = [(-x, x, y) for x, y in sushi[k:]]
heapq.heapify(ns)

for i in range(k - 1):
    if all([i == 1 for i in num_t.values()]):
        break
    for j in range(k - 1):
        d, t = heapq.heappop(s)
        if num_t[t] > 1:
            for l in range(n - k):
                dnm, dn, tn = heapq.heappop(ns)
                if not tn in num_t:
                    num_t[t] -= 1
                    num_d -= d
                    num_t[tn] = 1
                    num_d += dn
                    break
            break
        else:
            heapq.heappush(s, (d, t))
    ans = max(ans, num_d + len(num_t) ** 2)
print(ans)

Submission status: Runtime Error

Solution code:
import heapq

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

sushi.sort(reverse=True)
ans = 0
num_d = 0
num_t = dict()
for i in range(k):
    d, t = sushi[i]
    num_d += d
    if not t in num_t:
        num_t[t] = 1
    else:
        num_t[t] += 1
ans = max(ans, num_d + len(num_t) ** 2)
s = sushi[:k]
heapq.heapify(s)
ns = [(-x, x, y) for x, y in sushi[k:]]
heapq.heapify(ns)

for i in range(k - 1):
    if all([i == 1 for i in num_t.values()]):
        break
    for j in range(k - 1):
        d, t = heapq.heappop(s)
        if num_t[t] > 1:
            for l in range(n - k):
                dnm, dn, tn = heapq.heappop(ns)
                if not tn in num_t:
                    num_t[t] -= 1
                    num_d -= d
                    num_t[tn] = 1
                    num_d += dn
                    break
            break
        else:
            heapq.heappush(s, (d, t))
    ans = max(ans, num_d + len(num_t) ** 2)
print(ans)

Submission status: Runtime Error

Solution code:
import heapq

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():
        heapq.heappush(x, di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

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

print(ans)

Submission status: Runtime Error

Solution code:
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):
    if sushi[i][0] in t.keys():
        x.append(i)
        t[sushi[i][0]] += 1
    else:
        t[sushi[i][0]] = 1
    z += sushi[i][1]

y = list(t.keys())
ans = z + len(y) ** 2
k = 1
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in y:
        y.append(sushi[i][0])
        xi = x[-k]
        z += sushi[i][1] - sushi[xi][1]
        ans = max(ans, z + len(y) ** 2)
        k += 1

print(ans)

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

Submission status: Runtime Error

Solution code:
import copy
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 = {}
x = []
z = [0]
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
        t[sushi[i][0]] += 1
    else:
        t[sushi[i][0]] = 1
    z[0] += sushi[i][1]
y = []
yd = []
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in t.keys() and sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])

for k in range(1, len(x) + 1):
    xi = x[-k]
    z.append(z[-1] - sushi[xi][1] + sushi[y[k - 1]][1])

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

Submission status: Runtime Error

Solution code:
import copy
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 = []
z = [0]
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]
y = []
yd = []
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in t.keys() and 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):
    xi = x[-k]
    z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])

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

Submission status: Runtime Error

Solution code:
from heapq import heappop
from heapq import heappush
from heapq import heappushpop


FIN=10**10
N,K=map(int,input().split())

a=[list(map(int,input().split()))[::-1] for i in range(N) ]
a.sort(reverse=True)

l=[]
dic1={}
dic2={}

for i in range(K):
    p=a[i][1]
    if p in dic1:
        if a[i][0]>dic2[p]:
            heappush(l,dic2[p])
            dic2[p]=a[i][0]
        else: heappush(l,a[i][0])

    else: 
        dic1[p]=1
        dic2[p]=a[i][0]


s=sum(l)+sum(dic2.values())
x=len(dic1)
ans=s+x**2

for i in range(K,N):
    if not a[i][1] in dic1:
        dic1[a[i][1]]=1
        s+=a[i][0]-heappop(l)
        x+=1
        ans=max(ans,s+x**2) 

    if l==[]: break

print(ans)

Submission status: Runtime Error

Solution code:
# 最大価値が取れて種類数も増えるなら必ずとる
# 
# 種類数を固定すると：各種類のベストを見て貪欲に種類を決めているはず
# あとは種類数を更新できないものの中から貪欲
#
# 最大価値が等しい2種類：第2種の貪欲が使われだすなら両方取ってるはずなので優先度は同じでよい
#
# 種類数を増やしたときの変化を見る

import sys
input = sys.stdin.readline
from heapq import heappush,heappop

N,K = map(int,input().split())
kind_value = [[] for _ in range(N+1)]
for _ in range(N):
  t,d = map(int,input().split())
  kind_value[t].append(d)
for li in kind_value:
  li.sort()
  
first_type = [(li[-1],i) for i,li in enumerate(kind_value) if li]
first_type.sort(reverse = True)

kind = 0
value_type1 = 0
second_type = []
second_type_sum = 0
second_type_cnt = 0

answer = 0
for v,i in first_type:
  kind += 1
  value_type1 += v
  for d in kind_value[i][:-1]:
    heappush(second_type,d)
    second_type_sum += d
    second_type_cnt += 1
  while second_type_cnt > K-kind:
    d = heappop(second_type)
    second_type_cnt -= 1
    second_type_sum -= d
  value = kind*kind + value_type1 + second_type_sum
  if answer < value:
    answer = value

print(answer)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d

"""
import sys

from collections import Counter

def solve(N, K, sushi):
    sushi.sort(key=lambda x: x[1], reverse=True)
    pick = sushi[:K]
    pick_type = set([t for t, d in pick])

    cnt = Counter()
    for t, d in pick:
        cnt[t] += 1

    rem = dict()
    for t, d in sushi:
        if t not in pick_type:
            rem[t] = max(rem.get(t, 0), d)

    res = sorted([[k, v] for k, v in rem.items()], key=lambda x: x[1])

    manzoku = sum([d for t, d in pick])
    ans = [manzoku + len(pick_type)**2]
    while len(pick_type) < K:
        t, d = res.pop()
        pick_type.add(t)
        tt, dd = pick.pop()
        if cnt[tt] > 1:
            cnt[tt] -= 1
            manzoku = manzoku - dd + d
            ans.append(manzoku+ len(pick_type)**2)
    return max(ans)


def main(args):
    N, K = map(int, input().split())
    sushi = [[int(i) for i in input().split()] for _ in range(N)]
    ans = solve(N, K, sushi)
    print(ans)


if __name__ == '__main__':
    main(sys.argv[1:])

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d

"""
import sys

from collections import Counter

def solve(N, K, sushi):
    sushi.sort(key=lambda x: x[1], reverse=True)
    pick = sushi[:K]
    pick_type = set([t for t, d in pick])

    cnt = Counter()
    for t, d in pick:
        cnt[t] += 1

    rem = dict()
    for t, d in sushi:
        if t not in pick_type:
            rem[t] = max(rem.get(t, 0), d)

    res = sorted([[k, v] for k, v in rem.items()], key=lambda x: x[1])

    manzoku = sum([d for t, d in pick])
    ans = [manzoku + len(pick_type)**2]
    while res:
        t, d = res.pop()
        pick_type.add(t)
        tt, dd = pick.pop()
        if cnt[tt] > 1:
            cnt[tt] -= 1
            manzoku = manzoku - dd + d
            ans.append(manzoku+ len(pick_type)**2)
    return max(ans)


def main(args):
    N, K = map(int, input().split())
    sushi = [[int(i) for i in input().split()] for _ in range(N)]
    ans = solve(N, K, sushi)
    print(ans)


if __name__ == '__main__':
    main(sys.argv[1:])