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:
# ABC 116 D
import heapq

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

# 貪欲にとる.
keep = []  # 各種類のsushiでおいしさ最大のものを保存
exchange = []  # 取り替える用 (種類がかぶっているものが入る). priority queue
types = set() # keepに入っているsushiの種類
for i in range(K):
    d, t = sushi[i]
    if t not in types:
        keep.append(d)
        types.add(t)
    else:
        heapq.heappush(exchange, d)

basic_point_1 = sum(keep)
basic_point_2 = sum(exchange)
types_bonus = len(types)
ans = basic_point_1 + basic_point_2 + types_bonus**2

# 種類が増えるようにおいしさ最大のものから交換していく.
for i in range(K, N):
    if len(exchange) == 0:  # 取り替えると種類が減るので終了.
        break
    d, t = sushi[i]
    if t in types:
        continue
    types.add(t)
    basic_point_1 += d
    basic_point_2 -= heapq.heappop(exchange) # おいしさ最小のものと交換.
    types_bonus += 1
    ans = max(ans, basic_point_1 + basic_point_2 + types_bonus**2)

print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())

top = [0]*n # t毎に最高のdを格納するlist
sub = [] # それ以外のdを格納するlist
cst = [0] # cumsum top
css = [0] # cumsum sub

for _ in range(n):
    t, d = (int(i) for i in input().split())
    t -= 1
    if top[t] == 0:
        top[t] = d
    elif d > top[t]:
        sub.append(top[t])
        top[t] = d
    else:
        sub.append(d)

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

for i in range(n): # カテゴリ内最大のdが高い順に累積
    if top[i] != 0:
        cst.append(cst[-1] + top[i])

for i in range(len(sub)): # それ以外の全てのカテゴリの寿司のdが高い順に累積
    css.append(css[-1] + sub[i])

po = 0
for i in range(1, k+1):
    if len(cst)-1 >= i and len(css)-1 >= k-i:
        po = max(cst[i] + css[k-i] + i*i, po)

print(po)

Submission status: Accepted

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

from itertools import islice
import operator

N, K = map(int, input().split())
sushi = [tuple([i] + [int(x) for x in input().split()]) for i in range(N)]
sushi.sort(key=operator.itemgetter(2), reverse=True)

from collections import OrderedDict, deque, Counter
left = OrderedDict((x[0], x) for x in sushi)

base = 0
choosen = OrderedDict()
choosen_kind = set()
for x in sushi:
    i, t, _ = x
    if t not in choosen_kind:
        choosen[t] = x
        choosen_kind.add(t)
        del left[i]

choosen = list(islice(choosen.values(), K))
base = sum(d for _,_,d in choosen)
kinds = len(choosen)

cands = islice(left.values(), K)
fixed = islice(cands, K - len(choosen))
base += sum(d for _,_,d in fixed)

score = base + kinds**2

for cand in cands:
    if cand[2] <= choosen[-1][2]:
        break

    base_n = base - choosen[-1][2] + cand[2]
    kinds_n = kinds - 1
    
    score_n = base_n + kinds_n**2
    if score_n > score:
        choosen.pop()
        base, kinds, score = base_n, kinds_n, score_n
    else:
        break

print(score)

Submission status: Accepted

Solution code:
# coding:utf-8

import sys
import heapq
from collections import defaultdict, deque

INF = float('inf')
MOD = 10 ** 9 + 7

def LI(): return [int(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 input()


n, k = LI()
sushi = [(d, t) for t, d in (LI() for _ in range(n))]
sushi.sort()
sushi.reverse()

# sushiを美味しいものからk個とる
umami_max = []  # 寿司の種類数が減らないように最低一つは別で確保
umami_sub = []  # 種類数を増やす時に捨てる候補(priority_queueで管理)
used = set()
for i in range(k):
    d, t = sushi[i]
    if t not in used:
        umami_max.append(d)
        used.add(t)
    else:
        heapq.heappush(umami_sub, d)

umami_1 = sum(umami_max)
umami_2 = sum(umami_sub)
sushi_v = len(used)
ans = umami_1 + umami_2 + sushi_v ** 2

for i in range(k, n):
    if len(umami_sub) == 0:  # 捨てる寿司の候補がない場合，満足度はこれ以上大きくならない
        break
    d, t = sushi[i]
    if t in used:
        continue
    used.add(t)
    umami_1 += d
    umami_2 -= heapq.heappop(umami_sub)  # 一番美味しくない寿司を消す
    sushi_v += 1
    ans = max(ans, umami_1 + umami_2 + sushi_v ** 2)

print(ans)

Submission status: Accepted

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

from itertools import islice
import operator

N, K = map(int, input().split())
sushi = [tuple([i] + [int(x) for x in input().split()]) for i in range(N)]
sushi.sort(key=operator.itemgetter(2), reverse=True)

from collections import OrderedDict, deque, Counter
left = OrderedDict((x[0], x) for x in sushi)

base = 0
choosen = [None] * N
choosen_kind = set()
ci = 0
for x in sushi:
    i, t, _ = x
    if t not in choosen_kind:
        choosen[ci] = x
        ci += 1
        choosen_kind.add(t)
        del left[i]

choosen = choosen[:min(ci, K)]
base = sum(d for _,_,d in choosen)
kinds = len(choosen)

cands = islice(left.values(), K)
fixed = islice(cands, K - len(choosen))
base += sum(d for _,_,d in fixed)

score = base + kinds**2

for cand in cands:
    if cand[2] <= choosen[-1][2]:
        break

    base_n = base
    base_n -= choosen[-1][2]
    base_n += cand[2]
    kinds_n = kinds - 1
    
    score_n = base_n + kinds_n**2
    if score_n > score:
        choosen.pop()
        base, kinds, score = base_n, kinds_n, score_n
    else:
        break

print(score)

Submission status: Accepted

Solution code:
import numpy as np

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

# 美味しさでソート
taste_sorted_sushi = sorted(sushi, key=lambda x:x[1], reverse=True)

# 貪欲にK個選び分割する
selected_sushi = np.array(taste_sorted_sushi[:k], dtype=np.int64)

# 貪欲なK個の種類の分布を調べる
# K個の中で二つ以上あるものとそうでないものに分ける
S = set()
q = []
for s in selected_sushi:
    if s[0] in S:
        q.append(s[1])
    else:
        S.add(s[0])

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

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

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

print(ans)

Submission status: Accepted

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

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

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

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

Submission status: Accepted

Solution code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][1])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            ans = max(ans, num)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
X = []
for i in range(N):
    x, y = map(int, input().split())
    X.append([i, x, y])

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

p = 0
L = [0] * (N+1)
bonus = 0
for i in range(K):
    p += X[i][2]
    if L[X[i][1]] == 0:
        bonus += 1
        
    L[X[i][1]] += 1

p += bonus ** 2
ma = p
i = K-1
j = K
while True:
    while j < N and L[X[j][1]]:
        j += 1
    if j == N:
        break
    else:
        while i >= 0 and L[X[i][1]] <= 1:
            i -= 1
        if i < 0:
            break
        else:
            L[X[i][1]] -= 1
            L[X[j][1]] += 1
            p -= X[i][2]
            p += X[j][2]
            p += bonus * 2 + 1
            bonus += 1
            ma = max(ma, p)
            
    i -= 1
    j += 1
        
print(ma)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
X = []
for i in range(N):
    x, y = map(int, input().split())
    X.append([i, x, y])

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

p = 0
L = [0] * (N + 1)
bonus = 0
for i in range(K):
    p += X[i][2]
    if L[X[i][1]] == 0:
        bonus += 1

    L[X[i][1]] += 1

p += bonus ** 2
ma = p
i = K - 1
j = K
while True:
    while j < N and L[X[j][1]]:
        j += 1
    if j == N:
        break
    else:
        while i >= 0 and L[X[i][1]] <= 1:
            i -= 1
        if i < 0:
            break
        else:
            L[X[i][1]] -= 1
            L[X[j][1]] += 1
            p -= X[i][2]
            p += X[j][2]
            p += bonus * 2 + 1
            bonus += 1
            ma = max(ma, p)
            i -= 1

    j += 1

print(ma)

Submission status: Accepted

Solution code:
n, k = [int(i) for i in input().split()]
sushi = [[int(i) for i in input().split()] for _ in range(n)]
sushi.sort(key=lambda x: x[1], reverse=True)

#print(sushi)

num = 0
kind = set()
changable = []

for i in range(k):
    num += sushi[i][1]
    if sushi[i][0] in kind:
        changable.append(sushi[i][1])
    else:
        kind.add(sushi[i][0])

#print(kind)
#print(changable)

num += len(kind) ** 2
ans = num

for i in range(k, n):
    if not changable:
        break

    if sushi[i][0] not in kind:
        num += sushi[i][1] + 2 * len(kind) + 1 - changable.pop()
        kind.add(sushi[i][0])
        ans = max(ans, num)

print(ans)

Submission status: Accepted

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    if count < K:
        while 1:
            if count >= K :
                break
            if len(not_eat_list) <= 0:
                break
            i =not_eat_list.popleft()
            eat_list.appendleft(i)
            max_score += sushi[i][0]
            already_eat.add(sushi[i][1])
            count += 1



    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    if len(not_eat_list) > 0:
        for new_idx in not_eat_list:
            if len(eat_list) == 0 :
                break
            eated_idx = eat_list.pop()
            if sushi[eated_idx][1] in already_eat:
                already_eat.remove(sushi[eated_idx][1])
            already_eat.add(sushi[new_idx][1])

            new_kind = len(already_eat)
            now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
            now_kind = new_kind
            max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import numpy as np
N, K = map(int,input().split())
td = np.empty((N,2),dtype=np.int64)
for i in range(N):
    td[i] = list(map(int,input().split()))

t = td[:,0]-1
d = td[:,1]
order1=np.argsort(d)[::-1]
mins = len(set(t[order1[:K]]))

best = -np.ones(N,dtype=np.int64)
for i in range(N):
    if best[t[i]] < 0 or d[i] > d[best[t[i]]]:
        best[t[i]]=i

bestmenu = np.array(list(set(best) - set([-1])),dtype=np.int64)
othermenu = np.array(list(set(range(N)) - set(bestmenu)),dtype=np.int64)

maxs = min(len(bestmenu),K)
ans = 0
bests=0
sorted_d_best=np.sort(d[bestmenu])[::-1]
sorted_d_other=np.sort(d[othermenu])[::-1]
for s in range(mins,maxs+1):
    score = s*s + sorted_d_best[:s].sum() + sorted_d_other[:(K-s)].sum()
    if score > ans:
        bests=s
        ans = score

print(ans)

Submission status: Accepted

Solution code:
import sys
fin = sys.stdin.readline
from heapq import heappop, heappush
from collections import defaultdict


N, K = [int(elem) for elem in fin().split()]
sushi_info = sorted([tuple(int(elem) for elem in fin().split()) for _ in range(N)], key=lambda x: -x[1])


cur_eaten = []
# cur_rest = []
count = defaultdict(int)
for t, d in sushi_info[:K]:
    if count[t] >= 1:
        heappush(cur_eaten, (d, t))
    count[t] += 1


# for t, d in sushi_info[K:]:
#     heappush(cur_rest, (-d, t))

num_distinctive_kinds = sum(v >= 1 for v in count.values())
satisfactions = [sum(d for _, d in sushi_info[:K]) + num_distinctive_kinds**2]

cur_num_kinds = num_distinctive_kinds
for i in range(K, N):
    if count[sushi_info[i][0]] >= 1:
        continue
    if not cur_eaten:
        break
    ditched_d, ditched_t = heappop(cur_eaten)
    added_t, added_d = sushi_info[i]
    # added_d *= -1
    count[ditched_t] -= 1
    count[added_t] += 1
    cur_num_kinds += 1
    satisfactions.append(satisfactions[-1] - ditched_d + added_d + cur_num_kinds * 2 - 1)

print(max(satisfactions))

Submission status: Accepted

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

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

dict_td = set()
q = []
ans = 0
for i in range(k):
    if td[i][0] in dict_td:
        q.append(td[i])
    ans += td[i][1]
    dict_td.add(td[i][0])

ans += len(dict_td) ** 2

tmp_ans = ans
for i in range(k, n):
    if len(q) == 0:
        break
    if td[i][0] not in dict_td :
        minus = q.pop()
        tmp_ans += td[i][1] - minus[1]
        dict_len = len(dict_td)
        tmp_ans += (dict_len + 1) ** 2 - dict_len ** 2
        dict_td.add(td[i][0])
    ans = max(tmp_ans, ans)

print(ans)

Submission status: Accepted

Solution code:
N,K=map(int,input().split())

ate=[0 for i in range(N+1)]
sushi=[]

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

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

x=0
res=0

for s in sushi[:K]:
    if ate[s[0]]==0:
        x+=1
    ate[s[0]]+=1
    res+=s[1]

repl=K
delind=K-1

ans = res+x**2

while delind>=0 and repl<=N-1:
    f = False
    while delind>=0:
        t,d=sushi[delind]
        if ate[t]>1:
            ate[t]-=1
            res-=d
            delind-=1
            f=True
            break
        delind-=1

    while repl<=N-1:
        t,d=sushi[repl]
        if ate[t]==0:
            ate[t]+=1
            x+=1
            res+=d
            repl+=1
            f=f and True
            break
        repl+=1
    if f:
        ans=max(ans,res+x**2)

print(ans)

Submission status: Accepted

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

Submission status: Accepted

Solution code:
N, K = [int(_) for _ in input().split()]
TD = [[int(_) for _ in input().split()] for i in range(N)]

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

d_sum = 0
t_set = set([])
d_surplus = []

for t, d in TD_sorted[:K]:
	d_sum += d
	if not t in t_set:
		t_set.add(t)
	else:
		d_surplus.append(d)

a = d_sum + len(t_set) ** 2
answer = a

for t, d, in TD_sorted[K:]:
	if not d_surplus:
		break
	if not t in t_set:
		a += 2 * len(t_set) + 1
		a += d
		a -= d_surplus.pop()
		t_set.add(t)
		answer = max(answer, a)

print(answer)

Submission status: Accepted

Solution code:
if __name__ == '__main__':
    import sys
    input = sys.stdin.readline
    import collections
    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)
    ans_list = [0]
    reverse_td = list(zip(*td))
    c = collections.Counter(reverse_td[0][:k])
    kind = len(c)
    ans = sum(reverse_td[1][:k])+kind**2
    check = k-1
    flag=False
    for i in range(k,n):
        t = td[i][0]
        d = td[i][1]
        if t in c:
            pass
        else:
            c[t]=1
            while True:
                new_t = td[check][0]
                if c[new_t]==1:
                    check -=1
                    if check == -1:
                        flag =True
                        break
                else:
                    c[new_t]-=1
                    ans_list.append(ans_list[-1]+2*kind+1+d-td[check][1])
                    check-=1
                    kind+=1
                    break
            if flag:
                break
    print(ans+max(ans_list))

Submission status: Accepted

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], reverse=True)
num = 0
kind = set()
changable = []
for i in range(k):
	num += sushi[i][1]
	if sushi[i][0] in kind:
		changable.append(sushi[i][1])
	else:
		kind.add(sushi[i][0])

num += len(kind)**2
ans = num

for i in range(k, n):
	if not changable:
		break

	if sushi[i][0] not in kind:
		num += sushi[i][1] + 2*len(kind)+1 - changable.pop()
		kind.add(sushi[i][0])
		ans = max(ans, num)

print(ans)

Submission status: Accepted

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], reverse=True)
#print(sushi)
num = 0
kind = set()
changable = []
for i in range(k):
	num += sushi[i][1]
	if sushi[i][0] in kind:
		changable.append(sushi[i][1])
	else:
		kind.add(sushi[i][0])

num += len(kind)**2
#changable.reverse()
#print(changable)
ans = num

for i in range(k, n):
	if not changable:
		break

	if sushi[i][0] not in kind:
		num += sushi[i][1] + 2*len(kind)+1 - changable.pop()
		#print(num, kind)
		kind.add(sushi[i][0])
		ans = max(ans, num)


print(ans)

Submission status: Accepted

Solution code:
import itertools
N, K = [int(_) for _ in input().split()]
TD = sorted([[int(_) for _ in input().split()] for _ in range(N)],
            key=lambda x: -x[1])
neta = set()
dupl = []
ans = [0]
for t, d in TD[:K]:
    ans[0] += d
    if t in neta:
        dupl += [d]
    else:
        neta.add(t)
x = len(neta)
ans[0] += x**2
for t, d in TD[K:]:
    if not t in neta and dupl:
        neta.add(t)
        ans += [ans[-1] - x**2 + (x + 1)**2 - dupl.pop() + d]
        x += 1
print(max(ans))

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
from heapq import heappop, heappush
from collections import defaultdict
def inpl(): return list(map(int, input().split()))
N, K = inpl()
DT = sorted([inpl()[::-1] for _ in range(N)], reverse=True)


dic = defaultdict(list)
counter = defaultdict(int)
for d, t in DT:
    dic[t].append(d)
    counter[t] += 1
each_max = {k: max(v) for k, v in dic.items()}
oic = 0

all_ = set(dic.keys())
eat = defaultdict(list)
for i in range(K):
    d, t = DT[i]
    eat[t].append(d)
    oic += d
x = len(eat)
score = oic + x*x
heap = []

for k, v in eat.items():
    if len(v) >= 2:
        heappush(heap, (min(v), k))

not_eat = sorted([max(dic[k]) for k in all_ - set(eat.keys())])
while heap:
    if len(not_eat) == 0:
        break
    v, k = heappop(heap)
    oic -= v - not_eat.pop()
    x += 1
    _ = eat[k].pop()
    if len(eat[k]) >= 2:
        heappush(heap, (min(eat[k]), k))

    score = max(score, oic + x*x)

print(score)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop

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

# initial
res = 0
sum = 0
q = []
used = {}
for i in range(K):
  (d, t) = heappop(h)
  if t not in used:
    used[t] = 1
  else:
    heappush(q, (-d, t))
  sum += -d
res = len(used)*len(used)+sum

# swap minimum
while q and h:
  (d, t) = heappop(h)
  if t not in used:
    used[t] = 1
    m, _ = heappop(q)
    sum = sum - m - d
  res = max(res, len(used)*len(used) + sum)
  
print(res)

Submission status: Accepted

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

X.sort(reverse=True)

p,v = 0,set([])
pp = []
for i in range(K):
    p += X[i][0]
    if X[i][1] not in v:
        v.add(X[i][1])
    else:
        pp.append((X[i][0],X[i][1]))

Ans = p + len(v)**2
ans = Ans
for i in range(K,N):
    if X[i][1] not in v and len(pp)>0:
        temp = pp[-1]
        pp = pp[:len(pp)-1]
        ans = ans - temp[0] + X[i][0] - len(v)**2 + (len(v)+1)**2
        v.add(X[i][1])
        Ans = max(Ans,ans)

print(Ans)

Submission status: Accepted

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

ANS=0

check=[0]*(10**5+1)
NOCHANGE=[]
NETA=[]

for i in range(K):
    ANS+=TD[i][1]
    NETA.append(TD[i][0])


for i in range(N):
    if check[TD[i][0]]==0:
        NOCHANGE.append(i)
        check[TD[i][0]]=1

NOCHANSET=set(NOCHANGE)
NETA=len(set(NETA))



#NOWLIST=set(list(range(K)))
#NOWLIST=sorted(list(NOWLIST-NOCHANSET))



NOCHANGE=sorted(list(NOCHANSET))

ANS+=NETA**2

#print(NOCHANGE)
#print(NOWLIST)
#print(ANS)
#print(NETA)

L=len(NOCHANGE)

for j in range(L):
    if NOCHANGE[j]>K-1:
        break
else:
    j=L
#print(j)
i=K-1

SCORE=ANS
#print(SCORE)
while j<L:
    while i in NOCHANSET:
        i-=1

    if i<0:
        break


    SCORE+=TD[NOCHANGE[j]][1]-TD[i][1]-NETA**2+(NETA+1)**2
    if SCORE>ANS:
        ANS=SCORE
    #print(i,j,NETA,SCORE)
    NETA+=1
    j+=1
    i-=1

    #print(SCORE)


    #print(i,j)

print(ANS)

Submission status: Accepted

Solution code:
from collections import defaultdict as dd
N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for _ in range(N)]
oishisa_sorted = sorted(sushi, key=lambda x: -x[1])

kiso = 0
n = 0
netanum = dd(int)
ans = 0
revidx = K - 1
for neta, oishisa in oishisa_sorted:
  if n < K:
    kiso += oishisa
    netanum[neta] += 1
    n += 1
    if n == K:
      ans = kiso + len(netanum) ** 2
  else:
    if netanum[neta] == 0:
      while revidx >= 0:
        neta2, oishisa2 = oishisa_sorted[revidx]
        revidx -= 1
        if netanum[neta2] >= 2:
          kiso = kiso - oishisa2 + oishisa
          netanum[neta2] -= 1
          netanum[neta] += 1
          tmp = kiso + len(netanum) ** 2
          ans = ans if ans > tmp else tmp
          break
print(ans)

Submission status: Accepted

Solution code:
import sys, itertools
def input(): return sys.stdin.readline()
def inpl(): return [int(i) for i in input().split()]
N, K = inpl()
Sushi = []
Neta = set()
ans = 0
for _ in range(N):
    t, d = inpl()
    Sushi.append((d, t))
Sushi.sort(reverse = True)
Db = []
Nd = []
for i in range(K):
    d, t = Sushi[i]
    if t not in Neta:
        Neta.add(t)
        ans += d
    else:
        Db.append(d)
neta = len(Neta)
Db = [0] + list(itertools.accumulate(Db))
Db.reverse()
for i in range(N-K):
    d, t = Sushi[K+i]
    if t not in Neta:
        Neta.add(t)
        Nd.append(d)
Nd = [0] + list(itertools.accumulate(Nd))
print(ans + max([(neta + i)**2 + j + k for i, (j, k) in enumerate(zip(Db, Nd))]))

Submission status: Accepted

Solution code:
import sys, itertools
#def input(): return sys.stdin.readline()
def inpl(): return [int(i) for i in input().split()]
N, K = inpl()
Sushi = []
Neta = set()
ans = 0
for _ in range(N):
    t, d = inpl()
    Sushi.append((d, t))
Sushi.sort(reverse = True)
Db = []
Nd = []
for i in range(K):
    d, t = Sushi[i]
    if t not in Neta:
        Neta.add(t)
        ans += d
    else:
        Db.append(d)
neta = len(Neta)
Db = [0] + list(itertools.accumulate(Db))
Db.reverse()
for i in range(N-K):
    d, t = Sushi[K+i]
    if t not in Neta:
        Neta.add(t)
        Nd.append(d)
Nd = [0] + list(itertools.accumulate(Nd))
print(ans + max([(neta + i)**2 + j + k for i, (j, k) in enumerate(zip(Db, Nd))]))

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
a = [list(map(int,input().split())) for i in range(n)]
mita = [0]*n
a.sort(key = lambda x : -x[1])
ireta = [0]*n
tmp=[]
for c in range(n):
  i,j = a[c]
  if mita[i-1] == 0:
    mita[i-1] = 1
    ireta[c] = 1
    tmp.append([i-1,j])
  if len(tmp) >= k:
    break
c = 0
kazu = sum(mita)
while len(tmp) < k:
  if ireta[c]:
    c+=1
    continue
  ireta[c] = 1
  mita[a[c][0]-1] += 1
  tmp.append(a[c])
tmp.sort(key = lambda x: -x[1])
c = 0
kaku = []
flag = 0
for _ in range(n):
  b = tmp.pop()
  if mita[b[0]] >1:
    tmp.append(b)
    break
  while ireta[c]:
    c+=1
    if c >= n:
      flag = 1
      tmp.append(b)
      break
  if flag:
    break
  if kazu*kazu+b[1] < (kazu-1)**2+a[c][1]:
    kaku.append(a[c])
    ireta[c] = 1
    mita[a[c][0]-1] += 1
    kazu -= 1
  else:
    tmp.append(b)
    break
ans = kazu*kazu
for i,j in tmp:
  ans += j
for i,j in kaku:
  ans += j
print(ans)

Submission status: Accepted

Solution code:
N, K = map( int, input().split())
T = [(0,0)]*N
for i in range(N):
    t, d = map( int, input().split())
    T[i] = (d,t)
T.sort( key= None, reverse = True)
C = T[:K]
AC = T[K:]
kiso = 0
s = 0
S = [0]*(N+1)
for i in range(K):
    d, t = T[i]
    kiso += d
    if S[t] == 0:
        S[t] = 1
        s += 1
    else:
        S[t] += 1
ans = kiso + s**2
C.sort()
r = 0

for i in range(K):
    d, t = C[i]
    if S[t] >= 2:
        while r < N-K:
            if S[ AC[r][1]] == 0:
                S[t] -= 1
                S[ AC[r][1]] += 1
                s += 1
                kiso -= d
                kiso += AC[r][0]
                r += 1
                break
            r += 1
    ans = max( ans, kiso + s**2)
print( ans)

Submission status: Accepted

Solution code:
N, K = [int(_) for _ in input().split()]
TD = [[int(_) for _ in input().split()] for i in range(N)]

# print(TD)

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

#print(xs)

baset = 0
ts = set([])
ds = []

for t, d in xs[:K]:
    baset += d
    if not t in ts:
        ts.add(t)
    else:
        ds.append(d)

result = baset + len(ts) ** 2
r = result

#print(r, ts, ds)

for t, d in xs[K:]:
    if not ds:
        break
    if not t in ts:
        r += 2 * len(ts) + 1
        r += d - ds.pop()
        ts.add(t)
        result = max(result, r)
        #print(r, ts, ds)

#print(r, ts, ds)

print(result)

Submission status: Accepted

Solution code:
N, K = map( int, input().split())
T = [(0,0)]*N
for i in range(N):
    t, d = map( int, input().split())
    T[i] = (d,t)
T.sort( key= None, reverse = True)
C = T[:K]
AC = T[K:]
kiso = 0
s = 0
S = [0]*(N+1)
for i in range(K):
    d, t = T[i]
    kiso += d
    if S[t] == 0:
        S[t] = 1
        s += 1
    else:
        S[t] += 1
ans = kiso + s**2
C.sort()
r = 0

for i in range(K):
    d, t = C[i]
    if S[t] >= 2:
        while r < N-K:
            if S[ AC[r][1]] == 0:
                S[t] -= 1
                S[ AC[r][1]] += 1
                s += 1
                kiso -= d
                kiso += AC[r][0]
                r += 1
                break
            r += 1
    ans = max( ans, kiso + s**2)
print( ans)

Submission status: Accepted

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

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
top = [0]*n
sub = []
for i in range(n):
    t,d = map(int,input().split())
    t -=1
    if top[t]==0:
        top[t]=d
    else:
        if top[t]>=d:
            sub.append(d)
        else:
            sub.append(top[t])
            top[t]=d

top.sort()
top = top[::-1]
res = [0]
for i in range(n):
    if top[i]!=0:
        res.append(res[-1]+top[i])
    else:
        res.append(-1)
sub.sort()
sub = sub[::-1]
w = [0]
for i in range(len(sub)):
    w.append(w[-1]+sub[i])
tmp = 0
for i in range(1,k+1):
    if len(w)-1>=k-i and res[i]!=-1:
        tmp = max(res[i]+w[k-i]+i*i,tmp)

print(tmp)