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: Wrong Answer

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

tds.sort(key=lambda x:x[1])

arr = [0] * (n + 1)

for idx, td in enumerate(tds):
    t, d = td
    arr[t] = [d, n - idx - 1]

tds = [y for x, y in tds[::-1]]

base = 0
res = 0
kind = 0
for xy in arr:
    if not xy:
        continue
    x, y = xy
    kind += 1
    base += x
    if y < k:
        tds[y] = 0
    else:
        while tds[k - 1] == 0:
            k -= 1
            if k == 0:
                break
        if k > 0:
            k -= 1
        if k == 0:
            break
    res = max(res, kind ** 2 + base + sum(tds[:k]))

print(res)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop
from operator import itemgetter
import sys

input = sys.stdin.readline

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

e.sort(key=itemgetter(0), reverse=True)

b = e[:k]  # 暫定選択（おいしさ基礎ポイント降順取得）

used = [False] * n
hq = []

ans = 0  # 暫定選択（おいしさ基礎ポイント降順取得）
tn = 0  # 種類数
for d, t in b:
    if not used[t]:
        used[t] = True
        tn += 1
    else:
        heappush(hq, (d, t))
    ans += d
ans += pow(tn, 2)  # 種類ボーナス

c = e[k:]  # 取得しなかった予備

# print(b)
# print(c)
# print(hq)
# print(used)
s = ans
cur = 0
while hq:
    d, t = heappop(hq)
    s -= d
    changed = False
    while cur < n - k:
        cd, ct = c[cur]
        cur += 1
        if not used[ct]:
            changed = True
            s += cd
            s += tn * 2 + 1
            tn += 1
            if s > ans:
                ans = s
            break
    if not changed:
        break
print(ans)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

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

sushi = []

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

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

# 美味しいものから順にK個選ぶ
nowSushi = []
sumTaste = 0
numKind = 0
take = [0] * N  # 種類は高々N

for i in range(K) :
    t, d = sushi[i]
    nowSushi.append((t,d))
    sumTaste += d
    if take[t-1] == 0 :
        numKind += 1
    take[t-1] += 1

addSushi = []  # 種類を増やせる寿司
for i in range(K,N) :
    t, d = sushi[i]
    if take[t-1] > 0 :  # すでに選択済み
        continue
    addSushi.append((t,d))

L = min(len(addSushi), K)
ans = sumTaste + numKind**2

for i in range(L) :
    removeSushiT, removeSushiD = nowSushi.pop()
    if take[removeSushiT-1] <= 1 :
        continue

    addSushiT, addSushiD = addSushi[i]

    sumTaste = sumTaste - removeSushiD + addSushiD
    numKind = numKind + 1

    if (sumTaste + numKind**2) > ans :
        ans = sumTaste + numKind**2

print(ans)

Submission status: Wrong Answer

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: Wrong Answer

Solution code:
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
total = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    total -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, d)
ans = len(group)**2 + total

while td and td_:
    d, t = heapq.heappop(td)
    
    if t not in group:
        group.add(t)
        d_ = heapq.heappop(td_)
        total = total + d_ - d
        ans = max(ans, len(group)**2 + total)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
ans = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    ans -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, d)

length = len(group)
ans += length**2
n = ans

while td and td_:
    d, t = heapq.heappop(td)
    if t not in group:
        d_ = heapq.heappop(td_)
        group.add(t)
        n = n + d_ - d -length**2 + (length+1)**2
        length += 1
        ans = max(ans, n)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
ans = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    ans -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, d)

length = len(group)
n = ans

while td and td_:
    d, t = heapq.heappop(td)

    if t not in group:
        d_ = heapq.heappop(td_)
        group.add(t)
        n = n + d_ - d

print(max(ans+length**2, n+len(group)**2))

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(100001)
def input():
    return sys.stdin.readline()[:-1]

from heapq import *
from collections import defaultdict
N, K = map(int, input().split())
t_d = [list(map(int, input().split())) for i in range(N)]
t_d.sort(key = lambda x: -x[1])
dic = defaultdict(int)
mc, f, tmp, q1 = 0, [0 for _ in range(N+1)], 0, []
for t, d in t_d[:K]:
    if dic[t] == 0:
        mc += 1
    dic[t] += 1
    tmp += d
    heappush(q1, [d, t])
f[mc] = tmp
q2 = []
tmp = t_d[K:N]
for t, d in tmp[::-1]:
    q2.append([d, t])
# print(q1, q2)
c = mc
while q2:
    d, t = q2.pop()
    if dic[t] > 0:
        continue
    tmp = []
    while q1:
        _d, _t = heappop(q1)
        if dic[_t] > 1:
            dic[t] += 1
            dic[_t] -= 1
            c += 1
            for x in tmp:
                heappush(q1, x)
            heappush(q1, [d, t])
            break
        tmp.append([_d, _t])
    f[c] = f[c-1] + d - _d
    # print(q1, q2)
ans = 0
for c in range(mc, N):
    # print(c, f[c])
    if f[c] == 0:
        break
    ans = max(ans, f[c] + c**2)
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(100001)
def input():
    return sys.stdin.readline()[:-1]

from heapq import *
from collections import defaultdict
N, K = map(int, input().split())
t_d = [list(map(int, input().split())) for i in range(N)]
t_d.sort(key = lambda x: -x[1])
dic = defaultdict(int)
mc, f, tmp, q1 = 0, [0 for _ in range(N+1)], 0, []
for t, d in t_d[:K]:
    if dic[t] == 0:
        mc += 1
    dic[t] += 1
    tmp += d
    heappush(q1, [d, t])
f[mc] = tmp
q2 = []
tmp = t_d[K:N]
for t, d in tmp[::-1]:
    q2.append([d, t])
# print(q1, q2)
c = mc
while q2:
    d, t = q2.pop()
    if dic[t] > 0:
        continue
    tmp = []
    while q1:
        _d, _t = heappop(q1)
        if dic[_t] > 1:
            break
        tmp.append([_d, _t])
    dic[_t] -= 1
    dic[t] += 1
    c += 1
    for x in tmp:
        heappush(q1, x)
    heappush(q1, [-d, t])
    f[c] = f[c-1] + d - _d
    # print(q1, q2)
ans = 0
for c in range(mc, N):
    # print(c, f[c])
    if f[c] == 0:
        break
    ans = max(ans, f[c] + c**2)
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(100001)
def input():
    return sys.stdin.readline()[:-1]

from heapq import *
from collections import defaultdict
N, K = map(int, input().split())
t_d = [list(map(int, input().split())) for i in range(N)]
t_d.sort(key = lambda x: -x[1])
dic = defaultdict(int)
mc, f, tmp, q1 = 0, [0 for _ in range(N+1)], 0, []
for t, d in t_d[:K]:
    if dic[t] == 0:
        mc += 1
    dic[t] += 1
    tmp += d
    heappush(q1, [d, t])
f[mc] = tmp
q2 = []
for t, d in t_d[K:]:
    heappush(q2, [-d, t])
c = mc
while q2:
    d, t = heappop(q2)
    if dic[t] > 0:
        continue
    tmp = []
    while q1:
        _d, _t = heappop(q1)
        if dic[_t] > 1:
            break
        tmp.append([_d, _t])
    dic[_t] -= 1
    dic[t] += 1
    c += 1
    for x in tmp:
        heappush(q1, x)
    heappush(q1, [-d, t])
    f[c] = f[c-1] - d - _d
ans = 0
for c in range(mc, N):
    if f[c] == 0:
        break
    ans = max(ans, f[c] + c**2)
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(100001)
def input():
    return sys.stdin.readline()[:-1]

from heapq import *
from collections import defaultdict
N, K = map(int, input().split())
t_d = sorted([list(map(int, input().split())) for i in range(N)], key = lambda x : -x[1])
dic = defaultdict(int)
mc, f, tmp = 0, [0 for _ in range(N+1)], 0
for t, d in t_d[:K]:
    if dic[t] == 0:
        mc += 1
    dic[t] += 1
    tmp += d
f[mc] = tmp
q1 = t_d[:K]
q2 = t_d[K:]
heapify(q1)
heapify(q2)
c = mc
while q2:
    t, d = heappop(q2)
    if dic[t] > 0:
        continue
    tmp = []
    while q1:
        _t, _d = heappop(q1)
        if dic[_t] > 1:
            break
        tmp.append([_t, _d])
    dic[_t] -= 1
    dic[t] += 1
    c += 1
    for x in tmp:
        heappush(q1, x)
    heappush(q1, [t, d])
    f[c] = f[c-1] + d - _d
ans = 0
for c in range(mc, N):
    if f[c] == 0:
        break
    ans = max(ans, f[c] + c**2)
print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD += [[0, 0]]
TD.sort(key = lambda x: x[1], reverse = True)
TD.sort(key = lambda x: x[0]) #ネタの種類、おいしさを、ネタの種類でソート

lst = []
for i in range(1, N + 1):
    if TD[i-1][0] != TD[i][0]:
        lst += [[TD[i][1], 1]]
    else:
        lst += [[TD[i][1], 0]]

lst.sort(key = lambda x: x[1], reverse = True)
lst.sort(key = lambda x: x[0], reverse = True)
# print (lst)

total = 0
kind = 0
lst2 = []
for i in range(K):
    total += lst[i][0]
    if lst[i][1] == 1:
        kind += 1
    else:
        lst2 += [lst[i][0]]

ans = total + kind * kind

# print (lst2)
n = len(lst2)
j = K
for i in range(n):
    flag = True
    while flag:
        if j >= N:
            print (ans)
            exit()
        if lst[j][1] == 1: #ネタの種類が増えるものだけ考える
            total += lst[j][0]
            kind += 1
            total -= lst2[n-1-i]
            ans = max(ans, total + kind * kind)
            flag = False
        else:
            j += 1
print (ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD += [[0, 0]]
TD.sort(key = lambda x: x[1], reverse = True)
TD.sort(key = lambda x: x[0]) #ネタの種類、おいしさを、ネタの種類でソート

lst = []
for i in range(1, N + 1):
    if TD[i-1][0] != TD[i][0]:
        lst += [[TD[i][1], 1]]
    else:
        lst += [[TD[i][1], 0]]
lst.sort(reverse = True)
# print (lst)

total = 0
kind = 0
lst2 = []
for i in range(K):
    total += lst[i][0]
    if lst[i][1] == 1:
        kind += 1
    else:
        lst2 += [lst[i][0]]

ans = total + kind * kind

# print (lst2)
n = len(lst2)
j = K
for i in range(n):
    flag = True
    while flag:
        if j >= N:
            print (ans)
            exit()
        if lst[j][1] == 1: #ネタの種類が増えるものだけ考える
            total += lst[j][0]
            kind += 1
            total -= lst2[n-1-i]
            ans = max(ans, total + kind * kind)
            flag = False
        else:
            j += 1
print (ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD += [[0, 0]]
TD.sort(key = lambda x: x[0]) #ネタの種類、おいしさをネタの種類でソート

lst = []
for i in range(1, N + 1):
    if TD[i-1][0] != TD[i][0]:
        lst += [[TD[i][1], 1]]
    else:
        lst += [[TD[i][1], 0]]

lst.sort(reverse = True)
# print (lst)

total = 0
kind = 0
lst2 = []
for i in range(K):
    total += lst[i][0]
    if lst[i][1] == 1:
        kind += 1
    else:
        lst2 += [lst[i][0]]

ans = total + kind * kind

# print (lst2)
n = len(lst2)
j = K
for i in range(n):
    flag = True
    while flag:
        if j >= N:
            print (ans)
            exit()
        if lst[j][1] == 1:
            total += lst[j][0]
            kind += 1
            total -= lst2[n-1-i]
            ans = max(ans, total + kind * kind)
            flag = False
        else:
            j += 1
print (ans)

Submission status: Wrong Answer

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])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop()
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x
        # print("remove", remove_candidates)

    if i == k - 1:
        x = len(kinds.keys())
        ans_candidates[x] = umami + x * x
        remove_candidates.sort()
        # print(umami)
        # print(kinds, x)
        # print(remove_candidates)

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))

Submission status: Wrong Answer

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])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop(0)
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x

    if i == k - 1:
        x = len(kinds)
        ans_candidates[x] = umami + x * x
        remove_candidates.sort()
        print(umami)
        print(kinds)
        print(remove_candidates)

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))

Submission status: Wrong Answer

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])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    if i < k:
        if t in kinds:
            remove_candidates.append(d)
        kinds[t].append(i)
        umami += d
    elif t not in kinds and len(remove_candidates) > 0:
        kinds[t].append(i)
        x += 1
        d2 = remove_candidates.pop()
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x

    if i == k - 1:
        x = len(kinds)
        ans_candidates[x] = umami + x * x
        remove_candidates.sort()
        # print(umami)
        # print(kinds)
        # print(remove_candidates)

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))

Submission status: Wrong Answer

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])
# print(sushi)

kinds = defaultdict(list)
remove_candidates = []
umami = 0

ans_candidates = {}
for i in range(n):
    t, d = sushi[i]
    if i < k:
        if t in kinds.keys():
            remove_candidates.append([t, d])
        kinds[t].append(i)
        umami += d
    elif t not in kinds.keys() and len(remove_candidates) > 0:
        kinds[t].append(i)
        x = len(kinds)
        t2, d2 = remove_candidates.pop()
        umami = umami - d2 + d
        ans_candidates[x] = umami + x * x

    if i == k - 1:
        x = len(kinds)
        ans_candidates[x] = umami + x * x
        remove_candidates.sort(key=lambda x: x[1])
        # print(umami)
        # print(kinds)
        # print(remove_candidates)

# print()
# print(ans_candidates)
print(max(ans_candidates.values()))

Submission status: Wrong Answer

Solution code:
import heapq
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)
#'''
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
Inf = 10 ** 18
table = []
kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	heapq.heappush(table, (-d[i], t[i]))
	
eat = []
taste = 0

for i in range(K):
	tmp = heapq.heappop(table)
	flag = 1
	for j in range(len(eat)):
		if eat[j][1] == tmp[1]:
			flag = 0
			
	if flag == 1:
		kind += 1
	heapq.heappush(eat, (-tmp[0], tmp[1]))
	taste += -tmp[0]
	
tmp_score = taste + kind * kind
ans = tmp_score

can_pop = 1
can_push = 1

while can_pop * can_push == 1:
	can_pop = 0
	can_push = 1
	for i in range(len(eat)):
		tmp = heapq.heappop(eat)
		for j in range(len(eat)):
			if tmp[1] == eat[j][1]:
				can_pop = 1
				taste -= tmp[0]
				break
		
		if can_pop == 1:
			break
		else:
			heapq.heappush(eat, (Inf, tmp[1]))

	for i in range(len(table)):
		tmp = heapq.heappop(table)
		for j in range(len(eat)):
			if tmp[1] == eat[j][1]:
				can_push = 0
				break

		if can_push == 1:
			taste += -tmp[0]
			kind += 1
			heapq.heappush(eat, (-tmp[0], tmp[1]))
			break

	if can_pop * can_push == 1:
		tmp_score = taste + kind * kind
		if tmp_score > ans:
			ans = tmp_score

print(ans)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
sushi=[]
for i in range(n):
    sushi.append([])
for i in range(n):
    t,d=map(int,input().split())
    sushi[t-1].append(d)
Max=[]
Rest=[]
v=0
for i in sushi:
    if not i==[]:
        Max.append(max(i))
        i.remove(max(i))
        for j in i:
            Rest.append(j)
        v=v+1
Max.sort(reverse=True)
Rest.sort(reverse=True)
Max_s=[]
Rest_s=[0]
Sum=0
for i in Max:
    Sum=Sum+i
    Max_s.append(Sum)
Sum=0
l_r=0
for i in Rest:
    Sum=Sum+i
    Rest_s.append(Sum)
    l_r=l_r+1
ans=0
for i in range(1,min(k,v)+1):
    if k-i<=l_r:
        temp=Max_s[i-1]+Rest_s[k-i]+i**2
        if ans<temp:
            ans=temp
print(ans+1)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
sushi=[]
for i in range(n):
    sushi.append([])
for i in range(n):
    t,d=map(int,input().split())
    sushi[t-1].append(d)
Max=[]
Rest=[]
v=0
for i in sushi:
    if not i==[]:
        Max.append(max(i))
        i.remove(max(i))
        Rest=Rest+i
        v=v+1
Max.sort(reverse=True)
Rest.sort(reverse=True)
Max_s=[]
Rest_s=[0]
Sum=0
for i in Max:
    Sum=Sum+i
    Max_s.append(Sum)
Sum=0
l_r=0
for i in Rest:
    Sum=Sum+i
    Rest_s.append(Sum)
    l_r=l_r+1
ans=0
for i in range(1,min(k,v)+1):
    if k-i<=l_r:
        temp=Max_s[i-1]+Rest_s[k-i]+i**2
        if ans<temp:
            ans=temp
print(ans+1)

Submission status: Wrong Answer

Solution code:
from collections import deque
N, K = map(int, input().split())
dt = []
for _ in range(N):
    t, d = map(int, input().split())
    dt.append((d, t))
dt.sort(reverse=True)

s = set()
bests = []
rems = deque()
for d, t in dt:
    if t in s:
        rems.append(d)
    else:
        s.add(t)
        bests.append(d)
x = len(bests)
k = x
p = sum(bests) + x * x
while k < K:
    p += rems.popleft()
    k += 1
maxp = p
while len(rems) and rems[0] > bests[-1]:
    p -= bests.pop() + x * x
    x -= 1
    p += rems.popleft() + x * x
    maxp = max(maxp, p)
print(maxp)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for i in range(n)]
td = sorted(td, key=lambda x: x[1], reverse=True)
sushi = [[] for i in range(2)]
used = set()
for t, d in td:
    if t in used:
        sushi[0].append(d)
    else:
        sushi[1].append(d)
        used.add(t)
y = k - min(k, len(sushi[0]))
sushi = list(reversed(sushi[0][:k])) + sushi[1]
sum = sum(sushi[:k])
ans = sum + y ** 2
for i in range(len(sushi) - k):
    sum += sushi[k + i] - sushi[i]
    ans = max(ans, sum + (y + i + 1) ** 2)
print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for i in range(n)]
td = sorted(td, key=lambda x: x[1], reverse=True)
sushi = [[] for i in range(2)]
used = set()
for t, d in td:
    if t in used:
        sushi[0].append(d)
    else:
        sushi[1].append(d)
        used.add(t)
y = k - min(k, len(sushi[0]))
sushi = list(reversed(sushi[0])) + sushi[1]
sum = sum(sushi[:k])
ans = sum + y ** 2
for i in range(len(sushi) - k):
    sum += sushi[k + i] - sushi[i]
    ans = max(ans, sum + (y + i + 1) ** 2)
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
N, K = map(int, input().split())

sushi_array = []

sushi_count = [-1] * N
sushi_kind = 0
for i in range(N):
    t, d = map(int, input().split())
    sushi_array.append([t - 1, d])
    if sushi_count[t - 1] == -1:
        sushi_count[t - 1] = 0
        sushi_kind += 1

sushi_array = sorted(sushi_array, key=lambda x: -x[1])
# print(sushi_array, sushi_kind)

now_kind = 0
kiso = 0
for sushi in sushi_array[:K]:
    kiso += sushi[1]
    if sushi_count[sushi[0]] == 0:
        now_kind += 1
    sushi_count[sushi[0]] += 1

ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1

now_kind += 1
while now_kind <= sushi_kind:

    while add_index < N:
        sushi = sushi_array[add_index]
        # print(sushi)
        add_index += 1
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break

    else:
        break
    while remove_index >= 0:
        sushi = sushi_array[remove_index]
        # print(sushi)
        remove_index -= 1
        if sushi_count[sushi[0]] >= 1:
            kiso -= sushi[1]
            sushi_count[sushi[0]] -= 1
            break
    else:
        break
    # print(ans,kiso,now_kind * now_kind)
    ans = max(ans, kiso + now_kind * now_kind)
    now_kind += 1

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
N, K = map(int, input().split())

sushi_array = []

sushi_count = [-1] * N
sushi_kind = 0
for i in range(N):
    t, d = map(int, input().split())
    sushi_array.append([t - 1, d])
    if sushi_count[t - 1] == -1:
        sushi_count[t - 1] = 0
        sushi_kind += 1

sushi_array = sorted(sushi_array, key=lambda x: -x[1])
# print(sushi_array, sushi_kind)

now_kind = 0
kiso = 0
for sushi in sushi_array[:K]:
    kiso += sushi[1]
    if sushi_count[sushi[0]] == 0:
        now_kind += 1
    sushi_count[sushi[0]] += 1

ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1

now_kind += 1
while now_kind <= sushi_kind:

    while add_index < N:
        sushi = sushi_array[add_index]
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            add_index += 1
            break

    else:
        break
    while remove_index > 0:
        sushi = sushi_array[remove_index]
        if sushi_count[sushi[0]] >= 1:
            kiso -= sushi[1]
            sushi_count[sushi[0]] -= 1
            remove_index -= 1
            break
    else:
        break
    # print(kiso + now_kind * now_kind)
    ans = max(ans, kiso + now_kind * now_kind)
    now_kind += 1

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
N, K = map(int, input().split())

sushi_array = []

sushi_count = [-1] * N
sushi_kind = 0
for i in range(N):
    t, d = map(int, input().split())
    sushi_array.append([t - 1, d])
    if sushi_count[t - 1] == -1:
        sushi_count[t - 1] = 0
        sushi_kind += 1

sushi_array = sorted(sushi_array, key=lambda x: -x[1])
# print(sushi_array, sushi_kind)

now_kind = 0
kiso = 0
for sushi in sushi_array[:K]:
    kiso += sushi[1]
    if sushi_count[sushi[0]] == 0:
        now_kind += 1
    sushi_count[sushi[0]] += 1

ans = kiso + now_kind * now_kind

add_index = K
remove_index = K - 1

now_kind += 1
while now_kind <= sushi_kind:

    while add_index < N:
        sushi = sushi_array[add_index]
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break
        add_index += 1
    else:
        break
    while remove_index > 0:
        sushi = sushi_array[remove_index]
        if sushi_count[sushi[0]] >= 1:
            kiso -= sushi[1]
            sushi_count[sushi[0]] -= 1
            break
        remove_index -= 1
    else:
        break
    # print(kiso + now_kind * now_kind)
    ans = max(ans, kiso + now_kind * now_kind)
    now_kind += 1

print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
N, K = map(int, input().split())

sushi_dict = defaultdict()

for i in range(N):
    t, d = map(int, input().split())
    if t in sushi_dict:
        sushi_dict[t].append(d)
    else:
        sushi_dict[t] = [d]
# print(sushi_dict)

max_array = []
for key, value in sushi_dict.items():
    value.sort()
    max_array.append((key, value[-1]))

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

if K == 1:
    print(max_array[0][1] + 1)
    exit()

start_kind = min(K, len(max_array))
# print(max_array, start_kind)
ans = 0
for i in range(start_kind):
    ans += max_array[i][1]
ans += start_kind * start_kind
can_use_index = [max_array[i][0] for i in range(start_kind)]
can_use_num = []
for key, value in sushi_dict.items():
    if key not in can_use_index:
        continue
    for j in value[:-1]:
        can_use_num.append(j)

use_num_count = K - start_kind
ans += sum(can_use_num[:use_num_count])
# print(can_use_num)
# print(use_num_count, start_kind)
for i in range(start_kind - 1):
    if use_num_count + i >= len(can_use_num):
        break
    ans_new = can_use_num[use_num_count + i] + (start_kind - 1 - i) * (
        start_kind - 1 - i)
    ans_old = max_array[start_kind - 1 - i][1] + (start_kind - i) * (
        start_kind - i)
    # print(ans_new, ans_old)
    if ans_new > ans_old:
        ans += (ans_new - ans_old)
    else:
        break
print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
N, K = map(int, input().split())

sushi_dict = defaultdict()

for i in range(N):
    t, d = map(int, input().split())
    if t in sushi_dict:
        sushi_dict[t].append(d)
    else:
        sushi_dict[t] = [d]
# print(sushi_dict)

max_array = []
for key, value in sushi_dict.items():
    value.sort()
    max_array.append((key, value[-1]))

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

if K == 1:
    print(max_array[0][1] + 1)
    exit()

start_kind = min(K, len(max_array))
# print(max_array, start_kind)
ans = 0
for i in range(start_kind):
    ans += max_array[i][1]
ans += start_kind * start_kind
can_use_index = [max_array[i][0] for i in range(start_kind)]
can_use_num = []
for key, value in sushi_dict.items():
    if key not in can_use_index:
        continue
    for j in value[:-1]:
        can_use_num.append(j)

use_num_count = K - start_kind
ans += sum(can_use_num[:use_num_count])
# print(can_use_num)
# print(use_num_count, start_kind)
for i in range(start_kind):
    if use_num_count + i >= len(can_use_num):
        break
    ans_new = can_use_num[use_num_count + i] + (start_kind - 1 - i) * (
        start_kind - 1 - i)
    ans_old = max_array[start_kind - 1 - i][1] + (start_kind - i) * (
        start_kind - i)
    # print(ans_new, ans_old)
    if ans_new > ans_old:
        ans += (ans_new - ans_old)
    else:
        break
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
def LI(): return [int(x) for x in sys.stdin.readline().split()]

def main():
    N, K = LI()
    TD = []
    for _ in range(N):
        TD.append(LI())

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

    from collections import Counter
    import heapq

    counter = Counter()
    S = TD[:K]
    TD = TD[K:]
    queue = []
    for t, d in S:
        counter[t] += 1
        heapq.heappush(queue, (d, t))

    sum_ = sum([i[1] for i in S])
    maximum = sum_ + len(counter)**2
    while len(queue) > 0:
        d, t = heapq.heappop(queue)
        counter[t] -= 1
        if counter[t] == 0:
            continue
        sum_ -= d
        while len(TD) > 0:
            new_t, new_d = TD.pop(0)
            if counter[new_t] == 0:
                counter[new_t] += 1
                sum_ += new_d
                maximum = max(maximum, sum_ + len(counter)**2)
                break

    print(maximum)


main()