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

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
  if kind > K:
    break
  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: Accepted

Solution code:
# ABC 116 D - Various Sushi
N, K = map(int,input().split())

# S[k] := ネタkのおいしさのリスト
S = [[-float("inf")] for k in range(N)]
for k in range(N):
    t, d = map(int,input().split())
    S[t-1].append(d)

# 各種類のおいしさの最大を入れた配列がtop、それ以外がhoka
# それらを降順にソートして、k=1,2,..Kについて、topから大きい順にk個、hokaから同様にN-k個取ると、種類はk以上になる
# hokaからの要素で合計の種類がkより大きくなるばあい、その要素以上の大きさで同じ種類のものがtopに入っているため、
# あとでその種類からもtopを取る場合に、合計点がより大きい組み合わせになる
top = []
hoka = []
for k in range(N):
    S[k].sort()
    top.append(S[k].pop())
    hoka += S[k]

top.sort(reverse=True)
hoka.sort(reverse=True)

# 毎回sum(top[:k])とかやっていると間に合わないので累積和
# hokaから一つも取らないときは0を返すようにする
T = [0 for k in range(N+1)]
H = [0 for k in range(N+1)]
T[1], H[1] = top[0], hoka[0]
for k in range(2,N):
    T[k] = T[k-1] + top[k-1]
    H[k] = H[k-1] + hoka[k-1]

ans = 0
for k in range(1,K+1):
    ans = max(ans,T[k]+H[K-k]+k**2)
print(ans)

Submission status: Accepted

Solution code:
N, K = map(int,input().split())
S = [[-float("inf")] for k in range(N)]
for k in range(N):
    t, d = map(int,input().split())
    S[t-1].append(d)

top = []
hoka = []
for k in range(N):
    S[k].sort()
    top.append(S[k].pop())
    hoka += S[k]

top.sort(reverse=True)
hoka.sort(reverse=True)

T = [0 for k in range(N+1)]
H = [0 for k in range(N+1)]
T[1], H[1] = top[0], hoka[0]
for k in range(2,N):
    T[k] = T[k-1] + top[k-1]
    H[k] = H[k-1] + hoka[k-1]

ans = 0
for k in range(1,K+1):
    ans = max(ans,T[k]+H[K-k]+k**2)
print(ans)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
import sys
from collections import defaultdict
from collections import deque


def solve(N: int, K: int, T: "List[int]", D: "List[int]"):
    ml = []

    for t,d in zip(T,D):
        ml.append((t,d))
    ml.sort(key=lambda x: x[1],reverse=True)

    mli = ml[0:K]
    mlo = deque(ml[K:])
    mli.reverse()

    tdic = defaultdict(int)
    for i in mli:
        tdic[i[0]] += 1
    
    max_score = sum(map(lambda x:x[1],mli)) + len(tdic)**2
    prev_score = max_score
    for i,p in enumerate(mli):
        if tdic[p[0]] > 1:
            while len(mlo) > 0:
                q = mlo.popleft()
                if q[0] not in tdic:
                    tdic[p[0]] -= 1
                    tdic[q[0]] += 1
                    mli[i] = q
                    tdic_len  = len(tdic)
                    new_score = prev_score - p[1] + q[1] - (tdic_len-1)**2 + (tdic_len)**2
                    max_score = max(max_score,new_score)
                    prev_score = new_score
                    break

    print(max_score)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Submission status: Accepted

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

from collections import Counter

def solve(N, K, sushi):
    sushi.sort(key=lambda x: x[1], reverse=True)
    pick = sushi[:K]            #  おいしさポイントの高い方から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] #  現在の選択でのスコア、ここから種類を1つずつ増加させていきスコアを計算する
    while len(pick_type) < K and pick and res:
        t, d = res.pop()
        pick_type.add(t)
        tt, dd = pick.pop()
        while cnt[tt] == 1 and pick:
            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: Accepted

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]            #  おいしさポイントの高い方から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] #  現在の選択でのスコア、ここから種類を1つずつ増加させていきスコアを計算する
    while len(pick_type) < K and pick and res:
        t, d = res.pop()
        pick_type.add(t)
        tt, dd = pick.pop()
        while cnt[tt] == 1 and pick:
            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: Accepted

Solution code:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dt = sorted(zip(d, t), reverse=True)

    q0 = []

    acc = 0
    ts = 0
    p = [-1] * (N+1)
    for dd, tt in dt[:K]:
        acc += dd
        v = min(p[tt], dd)
        p[tt] = max(p[tt], dd)
        if 0 < v:
            heappush(q0, (v, tt))
        else:
            ts += 1

    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

    mmm = acc + ts * ts
    while 0 < len(q1) and 0 < len(q0):
        d1, t1 = heappop(q1)
        if 0 < p[t1]:
            continue
        acc -= d1
        p[t1] = -d1
        ts += 1

        d0, _ = heappop(q0)
        acc -= d0
        mmm = max(mmm, acc + ts * ts)

    return mmm


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    result = solve(N, K, t, d)
    print(result)

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
from collections import Counter
import sys
INF = 10 ** 18
def LI(): return [int(x) for x in sys.stdin.readline().split()]
def MI(): return map(int, sys.stdin.readline().split())


N, K = MI()
# 美味しさ降順寿司
sushi = sorted([LI() for _ in range(N)], key=lambda x:x[1], reverse=True)
# 美味しい寿司を貪欲に取る
delicious_sushi = sushi[:K]
# 余った寿司
soudemonai_sushi = sushi[K:]

syurui_counter = Counter()
deliciousness = 0

f = [-INF for _ in range(K+1)]

for x in delicious_sushi:
    syurui_counter[x[0]] += 1
    deliciousness += x[1]

f[len(syurui_counter)] = deliciousness

# print(delicious_sushi)
# print(soudemonai_sushi)
# print(syurui_counter)
# print("==========")

delicious_sushi = delicious_sushi[::-1]

i = 0
j = 0
while len(syurui_counter) < K and 0 < len(soudemonai_sushi):
    # print("len(syurui_counter) = " + str(len(syurui_counter)))
    while i < K:
        # print("i = " + str(i))
        if syurui_counter[delicious_sushi[i][0]] >= 2:
            # print(delicious_sushi[i])
            break
        i += 1
    while j < N - K:
        # print("j = " + str(j))
        if syurui_counter[soudemonai_sushi[j][0]] == 0:
            # print(soudemonai_sushi[j])
            break
        j += 1
    if j == N - K:
        break
    deliciousness -= delicious_sushi[i][1]
    deliciousness += soudemonai_sushi[j][1]
    syurui_counter[delicious_sushi[i][0]] -= 1
    syurui_counter[soudemonai_sushi[j][0]] += 1
    # print(syurui_counter)
    f[len(syurui_counter)] = deliciousness
    i += 1
    j += 1
    # print("==========")

# print(f)
for i in range(len(f)):
    f[i] += (i * i)
# print(f)
print(max(f))

Submission status: Accepted

Solution code:
"""
漸化式的に解く
一般性を失わない範囲をちゃんと考える。
"置き換え可能"というのをどの部分で持たせられるかを考える
今回の場合は、リスト作りの際に作ったので十分
"""

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] #変更できるもの(2個以上あるもの)



cnt = 0
for t,d in pairs[:K]: #K個上から順番に取ってみる
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

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

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    # 交換できるものがなくなったら終わり
    if len(change) == 0:
        break
    # 同じものは交換しない
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)
# 集計
print(ans)

Submission status: Accepted

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

type = set()
L = []
Sum = 0
for x in td[:k]:
    Sum += x[1]
    if x[0] not in type:
        type.add(x[0])
    else:
        L.append(x[1])

L = L[::-1]
type_cnt = len(type)
ans = Sum + type_cnt ** 2
for x in td[k:]:
    if len(L)==0:break
    if x[0] not in type:
        type.add(x[0])
        type_cnt += 1
        Sum = Sum - L.pop(0) + x[1]
        ans = max(ans, Sum + type_cnt ** 2)
print(ans)

Submission status: Accepted

Solution code:
N,K=[int(a) for a in input().split(" ")]
TD = [[int(a) for a in input().split(" ")] for i in range(N)]
TD = sorted(TD, key=lambda arr: arr[1])[::-1]

types = set()
stack = []
ans, totald, i = 0, 0, 0
for i in range(K):
    t, d = TD[i]
    if t in types:
        stack.append(d)
    totald += d
    types.add(t)
ans = totald+len(types)**2

while len(stack)>0:
    while i<N and TD[i][0] in types:
        i+=1
    if i>=N: break
    t1, d1 = TD[i]
    lt = len(types)
    d2 = stack.pop()
    totald -= (d2-d1)
    types.add(t1)
    i+=1
    ans = max(ans, totald+len(types)**2)
print(ans)

Submission status: Accepted

Solution code:
n,k = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for _ in range(n)]

###dが大きい順にソート
td.sort(key = lambda x: x[1],reverse=True)

###各tでdが最大のものを格納
fix_eated_d=[]
###各tでdが最大でないものを格納。tの種類を増やす際の削除候補
change_eated_d=[]

eated_t = set()
d_point = 0

###まずdを最大化
for t,d in td[:k]:
  
  ###すでにtが選ばれているなら最大のdではないのでchangeへ
  if t in eated_t:
    change_eated_d.append(d)
    
  ###まだtが選ばれていないなら最大のdだからfixへ
  else:
    fix_eated_d.append(d)
    eated_t.add(t)
    
  d_point += d
  
point = len(eated_t)**2 + d_point

###tの種類を増やす
for t,d in td[k:]:
  ###削除候補がなくなれば終了
  if not change_eated_d:
    break
  
  ###すでにtが選ばれているならskip
  if t in eated_t:
    continue
    
  ###まだtが選ばれていないなら追加、change中dが最小のものを削除
  eated_t.add(t)
  d_point += d - change_eated_d.pop()
  point = max(point, len(eated_t)**2 + d_point)
  
print(point)

Submission status: Accepted

Solution code:
N, K = list(map(int, input().split()))
td = sorted([list(map(int, input().split())) for _ in range(N)],key=lambda x:-x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += td[i][1]
    if td[i][0] in s:
        q.append(td[i][1])
    s.add(td[i][0])
ans += len(s)**2
k1 = K
ans1 = ans
while k1 < N and q:
    if td[k1][0] not in s:
        ans1 += td[k1][1] - q.pop() - len(s)**2 + (len(s)+1)**2
        ans = max(ans,ans1)
        s.add(td[k1][0])
    k1 += 1
print(ans)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
import sys
from collections import defaultdict
from itertools import accumulate
import heapq

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    sushi = defaultdict(list)
    for i in range(N):
        sushi[t[i]].append(d[i])

    A = [] # 各ネタでいちばんおいしい寿司のおいしさ
    B = [] # 各ネタで2番目以降においしい寿司のおいしさリスト
    
    for v in sushi.values():
        v.sort(reverse=True)
        A.append(v[0])
        B.extend(v[1:])
    A.sort(reverse=True)
    B.sort(reverse=True)

    X = [0] + [x for x in accumulate(A)] # X[i]: Aからうまい順にi個選んだ時のおいしさ合計
    Y = [0] + [y for y in accumulate(B)] # Y[j]: Bからうまい順にj個選んだ時のおいしさ合計

    res = 0
    for i in range(1, min(len(X), K + 1)):
        if K - i < len(Y):
            res = max(res, X[i] + Y[K - i] + i ** 2) 
            # B[:i]に含まれるネタが、A[i+1:]に含まれていた場合、実際のネタ数はiより多いが、
            # その場合はiを多くとったほうが、より美味しさが大きくなる。
            # そのためres[i]が答えになりえないため、無視して良い。
    print(res)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import sys
from collections import defaultdict
from itertools import accumulate

input = sys.stdin.readline

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

X, Y = [], []
for v in sushi.values():
    v = sorted(v, reverse=True)
    X.append(v[0])
    if len(v) > 1:
        Y += v[1:]
X = sorted(X, reverse=True)
Y = sorted(Y, reverse=True)

X = [0] + [x for x in accumulate(X)]
Y = [0] + [x for x in accumulate(Y)]

ans = 0
for i in range(1, min(len(X), K + 1)):
    j = K - i
    if j < len(Y):
        ans = max(ans, X[i] + Y[j] + i * i)
print(ans)

Submission status: Accepted

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

head = []
no_head = []
is_used = [False]*(N+1)
for t,d in td[:K]:
    if is_used[t]:
        no_head.append(d)
    else:
        head.append(d)
        is_used[t] = True

unused_head = []
for t,d in td[K:]:
    if not is_used[t]:
        unused_head.append(d)
        is_used[t] = True
        
n = min(len(no_head),len(unused_head))
sum_head = sum(head)
sum_no_head = sum(no_head)
sum_unused_head = 0
l_head = len(head)
l_no_head = len(no_head)
ans = sum_head + sum_no_head + l_head ** 2
for i in range(n):
    sum_no_head -= no_head[-1-i]
    sum_unused_head += unused_head[i]
    ans = max(ans,sum_head+sum_no_head+sum_unused_head + (l_head+i+1)**2)
print(ans)

Submission status: Accepted

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

typ = set()
ans = 0
cur = 0
rem = []

for t,d in sushi:
	if t in typ:
		if K > 0:
			K -= 1
			cur += d
			rem.append(d)
	else:
		typ.add(t)
		cur += d
		if K > 0:
			K -= 1
		elif len(rem):
			cur -= rem.pop()
		else:
			break

	ans = max(ans, cur + len(typ)**2)

print(ans)

Submission status: Accepted

Solution code:
import heapq

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

td = []
for _ in range(N):
    t, d = map(int, input().split())
    td.append((t, d))

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

g1 = [0] * (N + 1)
g0 = []
sumt = 0
for i in range(K):
    t, d = td[i]
    if g1[t] == 0:
        g1[t] = d
        sumt += 1
    elif g1[t] < d:
        heapq.heappush(g0, g1[t])
        g1[t] = d
    else:
        heapq.heappush(g0, d)

sumg1 = sum(g1)
sumg0 = sum(g0)
ans = sumg1 + sumg0 + sumt * sumt
for i in range(K, N):
    t, d = td[i]
    if g1[t] == 0 and len(g0) > 0:
        g1[t] = d
        sumg1 += d
        sumg0 -= heapq.heappop(g0)
        sumt += 1
        newans = sumg1 + sumg0 + sumt * sumt
        ans = max(ans, newans)

print(ans)

Submission status: Accepted

Solution code:
import heapq

N, K = map(int, input().split())
heap = []
types = [0 for _ in range(N)]

class Sushi():
    def __init__(self, neta, aji):
        self.neta = neta - 1
        self.aji = aji

    def __str__(self):
        return "Sushi{0}: {1}".format(self.neta, self.aji)

    def reverse(self):
        self.aji = -self.aji

    def __lt__(self, other):
        return self.aji < other.aji

for _ in range(N):
    t, d = map(int, input().split())
    heap.append(Sushi(t, -d))
heapq.heapify(heap)

variety = 0
basic = 0
choices = []
for _ in range(K):
    sushi = heapq.heappop(heap)
    sushi.reverse()
    basic += sushi.aji
    neta = sushi.neta
    if types[neta] == 0:
        variety += 1
    types[neta] += 1
    choices.append(sushi)
heapq.heapify(choices)

satisfaction = basic + variety * variety
top_satisfaction = satisfaction
while True:
    drop_sushi = None
    while len(choices) > 0:
        choiced_sushi = heapq.heappop(choices)
        if types[choiced_sushi.neta] > 1:
            drop_sushi = choiced_sushi
            types[choiced_sushi.neta] -= 1
            break
    add_sushi = None
    while len(heap) > 0:
        new_sushi = heapq.heappop(heap)
        new_sushi.reverse()
        if types[new_sushi.neta] == 0:
            types[new_sushi.neta] += 1
            add_sushi = new_sushi
            break
    if drop_sushi is None or add_sushi is None:
        break
    else:
        variety += 1
        basic += add_sushi.aji - drop_sushi.aji
        satisfaction = basic + variety * variety
        if satisfaction > top_satisfaction:
            top_satisfaction = satisfaction
print(top_satisfaction)

Submission status: Accepted

Solution code:
N,K=map(int, input().split())
td=[list(map(int,input().split())) for i in range(N)]
best=[0]*(N+1)
other=[]
from heapq import heappop,heappush
td.sort(key=lambda x: x[1],reverse=True)
shurui=0
for i in range(K):
    if best[td[i][0]]==0:
        best[td[i][0]]=td[i][1]
        shurui+=1
    elif best[td[i][0]]<td[i][1]:
        heappush(other,best[td[i][0]])
        best[td[i][0]]=td[i][1]
    else:
        heappush(other,td[i][1])

sumbest=sum(best)
sumother=sum(other)
ans=sumbest+sumother+shurui**2

for i in range(K,N):
    if best[td[i][0]]==0 and len(other)>0:
        best[td[i][0]]=td[i][1]
        shurui+=1
        sumbest+=td[i][1]
        sumother-=heappop(other)
        ans=max(ans,sumbest+sumother+shurui**2)

print(ans)

Submission status: Accepted

Solution code:
N , K = map(int, input().split())
sushi = [list(map(int, input().split())) for i in range(N)]
top = [-10**10]*(N+1)
sub = [0]
cnt = 0
for t, d in sushi:
    if top[t] < d:
        sub.append(top[t])
        top[t] = d
    else:
        sub.append(d)
top.sort(reverse = True)
sub.sort(reverse = True)
top_dp = [0] *(N+1)
sub_dp = [0] * (N+2)
top_dp[1] = top[0]
sub_dp[1] = sub[0]
for i in range(1, N):
    top_dp[i+1] = top_dp[i] + top[i]
for i in range(1, len(sub)):
    sub_dp[i+1] = sub_dp[i] + sub[i]
dp = [0] * (K+1)
for i in range(1, K+1):
    dp[i] = top_dp[i] + sub_dp[K-i] + i*i
print(max(dp))

Submission status: Accepted

Solution code:
nk = input().split()
n = int(nk[0])
k = int(nk[1])

dtlist = []
for i in range(n):
    td = input().split()
    t = int(td[0])
    d = int(td[1])
    dtlist.append((d, t))
dtlist.sort(reverse=True)


t_to_maxidx = {}
for idx, dt in enumerate(dtlist):
    d, t = dt
    if t not in t_to_maxidx:
        t_to_maxidx[t] = idx


selected = []
count = 0
tcount = {}

candidate = [(dtlist[idx][0], idx) for t, idx in t_to_maxidx.items()]
candidate.sort(reverse=True)
for d, idx in candidate:
    if count >= k:
        break
    d, t = dtlist[idx]
    selected.append(dtlist[idx])
    count += 1
    tcount[t] = 1 if t not in tcount else tcount[t] + 1
    dtlist[idx] = None

new_dtlist = []
for dt in dtlist:
    if dt is not None:
        new_dtlist.append(dt)
dtlist = new_dtlist


if k - count > 0:
    for i in range(k - count):
        dt = dtlist.pop(0)
        selected.append(dt)
        d, t = dt
        tcount[t] = 1 if t not in tcount else tcount[t] + 1

selected.sort()

score = 0
for d, t in selected:
    score += d

len_tcount = len(tcount)
score += len_tcount * len_tcount

for i in range(0, 2**32):
    next_score = score
    try:
        d_new, t_new = dtlist.pop(0)
        d_old, t_old = selected.pop(0)
    except IndexError:
        print(score)
        break
    next_score += d_new
    next_score -= d_old

    tcount[t_old] -= 1
    if tcount[t_old] <= 1:
        next_score -= len_tcount * len_tcount
        len_tcount -= 1
        next_score += len_tcount * len_tcount
        del tcount[t_old]

    if t_new not in tcount:
        next_score -= len_tcount * len_tcount
        len_tcount += 1
        next_score += len_tcount * len_tcount
        tcount[t_new] = 1
    else:
        tcount[t_new] += 1

    if next_score <= score:
        print(score)
        break
    score = next_score

Submission status: Accepted

Solution code:
import heapq
from collections import defaultdict

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

sushi = []
for _ in range(N):
    kind, value = map(int, input().split())
    sushi.append((value,kind))
sushi.sort()


choice = []
heapq.heapify(choice)

current_sum = 0
kinds = set()
count_nitems_bykind = defaultdict(int)

# pick K items with highest value
for _ in range(K):
    item = sushi.pop()
    count_nitems_bykind[item[1]] += 1
    kinds.add(item[1])
    current_sum += item[0]
    heapq.heappush(choice, item)

# initialize a dict to store max value that can be achieved by selecting x kinds of sushi
# key: x, value: max value
max_value = {len(kinds):current_sum}


def replace():
    '''
    Returns a next sushi that is new in kind and has maximum value
    '''

    while sushi:
        candidate = sushi.pop()
        if candidate[1] not in kinds:
            kinds.add(candidate[1])
            return candidate

    return None

while sushi and choice:
    # pick a sushi that has min value
    sushi_min = heapq.heappop(choice)

    # check if that sushi is a duplicate then if yes replace with a new sushi whose kind is new and value is max among such
    if count_nitems_bykind[sushi_min[1]] > 1:

        # pick a candidate to replace
        sushi_new = replace()
        if sushi_new is None:
            break

        count_nitems_bykind[sushi_min[1]] -= 1
        heapq.heappush(choice, sushi_new)


        # add a new result to dict
        current_sum = current_sum - sushi_min[0] + sushi_new[0]
        max_value[len(kinds)] = current_sum

    if len(kinds) == K: # already iterated enough to find answer
        break

ans = -1

for key in max_value:
    ans = max(ans, max_value[key] + key**2)

print(ans)

Submission status: Accepted

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

Submission status: Accepted

Solution code:
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([])
    base = 0
    for _c in choice:
        base += _c[1]
        if _c[0] in variety:
            r.append(_c[1])
        else:
            variety.add(_c[0])
    base += len(variety)**2
    i = k
    ans = [base]
    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)])))

Submission status: Accepted

Solution code:
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]
    count = [0] * n
    for _c in choice:
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_remove, i_append = k - 1, k
    while len(variety) < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        while count[td[i_remove][0] - 1] == 1:
            i_remove -= 1
        variety.add(td[i_append][0])
        count[td[i_remove][0] - 1] -= 1
        ans.append(ans[-1] - td[i_remove][1] + td[i_append][1] + 2 * len(variety) - 1)
        i_remove -= 1
        i_append += 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)])))

Submission status: Accepted

Solution code:
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]
    count = [0] * n
    r = []
    for _c in choice:
        if count[_c[0] - 1] > 0:
            r.append(_c[1])
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_append = k
    while len(variety) < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        variety.add(td[i_append][0])
        ans.append(ans[-1] - r.pop() + td[i_append][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)])))

Submission status: Accepted

Solution code:
from operator import itemgetter
from collections import Counter

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

sushi = sorted(sushi, key=itemgetter(1), reverse=True)

#
# kind_counter = Counter([s[0] for s in ate_list])
# kind_cnt = len(list(kind_counter))

##
# すでにでてきたかどうかがわかれば良いので Counter を使う必要はない
# (数える必要ない)
##

kind = set()
can_converts = []  # 変更しても良いよリスト
max_pt = 0

for i in range(K):
    if sushi[i][0] in kind:
        can_converts.append(sushi[i])
    else:
        kind.add(sushi[i][0])
    max_pt += sushi[i][1]
max_pt += len(kind) ** 2

# 食べた種類を増やしていくフェーズ
curr_pt = max_pt

for s in sushi[K:]:
    if not can_converts: break
    if s[0] not in kind:
        target = can_converts.pop(-1)
        curr_pt = curr_pt - target[1] + s[1] - len(kind) ** 2 + (len(kind) + 1) ** 2
        max_pt = max(max_pt,curr_pt)

        kind.add(s[0])

print(max_pt)

# kind_cnt_seq = kind_cnt ** 2
# max_pt = sum([s[1] for s in ate_list]) + kind_cnt_seq
#
# if K >= N:
#     print(max_pt)
#     exit()
#
# waiting_sushi = []
# wait_kind = []
# for s in sushi[K:]:
#     if kind_counter[s[0]]:  # すでに食べてる種類
#         continue
#
#     if not s[0] in wait_kind:
#         waiting_sushi.append(s)
#         wait_kind.append(s[0])
#
# # 食べた種類を増やしていくフェーズ
# idx = K - 1
# curr_pt = max_pt
# for s in waiting_sushi:
#     target = None
#     for i in range(idx, -1, -1):
#         # 1個しかないものは交換しない
#         if kind_counter[ate_list[i][0]] > 1:
#             target = ate_list[i]
#             idx = i
#             break
#
#     if not target:
#         break