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

N, K = map(int, input().split())
Items = [tuple(map(int, input().split())) for _ in range(N)]
Items.sort(reverse=True, key=lambda x: x[1])
Kinds = defaultdict(int)
ans = 0
k = 0
Candidates = []
for i, (t, d) in enumerate(Items):
    if i < K:
        ans += d
        if t not in Kinds:
            k += 1
        else:
            Candidates.append(d)
        Kinds[t] += 1
        if i == K - 1:
            ans += k ** 2
            a = ans
    else:
        if t in Kinds:
            continue
        if not Candidates:
            break
        else:
            x = Candidates.pop()
            a -= x
            a -= k ** 2
            k += 1
            a += d
            a += k ** 2
            ans = max(ans, a)
print(ans)

Submission status: Wrong Answer

Solution code:
from heapq import heapify, heappop, heappush
from collections import deque

n,k = map(int, input().split())
dtl = []
# t_dmax = [-1]*(n+1)
for _ in range(n):
    t,d = map(int, input().split())
    dtl.append((d,t))


dtl.sort(reverse=True)
dtq = deque(dtl)

dt_remq = deque([])
already_t = [False]*(n+1)
tll = [ [] for _ in range(n+1)]
cnt = 0
oisi = 0
kl = []
while dtq:
    d,t = dtq.popleft()
    if k <= cnt or already_t[t]: 
        dt_remq.append((d,t))
    else:
        already_t[t] = True
        kl.append((d,t))
        tll[t].append(d)
        cnt += 1
        oisi += d


rem = k-cnt
# if rem == 0:
#     ans = oisi + k*k
#     print(ans)
#     exit()

more_twos = []
for _ in range(rem):
    d,t = dt_remq.popleft()
    tll[t].append(d)
    oisi += d

ones = []
for tl in tll:
    if len(tl) == 1: ones.append(tl[0])

ones.sort()
onesq = deque(ones)

cuur_v = oisi + cnt*cnt
# print(cuur_v)
# print(kl)
while dt_remq and ones:
    d,t = dt_remq.popleft()
    oisi_diff = d-ones[0]
    bonus_diff = cnt*cnt-(cnt-1)*(cnt-1)
    if oisi_diff > bonus_diff:
        cuur_v += (oisi_diff-bonus_diff)
        cnt -= 1
    else:
        break

print(cuur_v)

Submission status: Wrong Answer

Solution code:
import sys
from collections import deque
from heapq import heappop, heappush

sys.setrecursionlimit(10 ** 7)
input = sys.stdin.readline
f_inf = float('inf')
mod = 10 ** 9 + 7


def resolve():
    n, k = map(int, input().split())
    LD = [[] for _ in range(n)]
    for _ in range(n):
        t, d = map(int, input().split())
        LD[t - 1].append([d, 0] if len(LD[t - 1]) else [d, 1])

    ALL = []
    for ld in LD:
        ALL.extend(ld)
    ALL.sort(reverse=True)
    que = deque(ALL)

    x = y = 0
    zero, one = [], []
    for _ in range(k):
        d, ex = que.popleft()
        x += d
        y += ex
        heappush(zero, d) if ex == 0 else heappush(one, d)

    res = x + y ** 2
    while que:
        d, ex = que.popleft()
        if ex:
            if zero:
                x += d - heappop(zero)
                y += 1
                res = max(res, x + y ** 2)
            else:
                break
    print(res)


if __name__ == '__main__':
    resolve()

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop


N, K, *td = map(int, open(0).read().split())
sushis = [(d, t) for t, d in zip(*[iter(td)] * 2)]
sushis.sort(reverse=True)

pq = []
cnt = [0] * (N + 1)
base_point = 0
for i in range(K):
    s = sushis[i]
    cnt[s[1]] += 1
    base_point += s[0]
    heappush(pq, s)
x = len([c for c in cnt if c > 0])

ans = base_point + x * x
for i in range(K, N):
    s = sushis[i]
    if cnt[s[1]] > 0:
        continue
    elif cnt[pq[0][1]] < 2:
        continue
    # remove
    d, t = heappop(pq)
    cnt[t] -= 1
    base_point -= d

    # add
    x += 1
    cnt[s[1]] += 1
    base_point += s[0]
    ans = max(ans, base_point + x * x)
    if not pq:
        break

print(ans)

Submission status: Wrong Answer

Solution code:
N,K=list(map(int,input().split()))
l=[tuple(map(int,input().split())) for i in range(N)]
l.sort(key=lambda x:x[1],reverse=True)
import heapq
print(l)
s_set=set()
length=0
var=0
change_l=[]
for i in range(K):
   length=len(s_set)
   s,n=l[i]
   s_set.add(s)
   if length!=len(s_set):
      var+=n
   else:
      change_l.append(n)
heapq.heapify(change_l)
ans=var+len(s_set)**2+sum(change_l)
for i in range(K,N):
   s,n=l[i]
   if s not in s_set and change_l!=[]:
      heapq.heappop(change_l)
      s_set.add(s)
      var+=n
      ans=max(ans,var+len(s_set)**2+sum(change_l))
print(ans)

Submission status: Wrong Answer

Solution code:
from heapq import heapify, heappop
from sys import stdin
def input():
    return stdin.readline().strip()

def main():
    n, k = map(int, input().split())

    dt = [(0, 0)] * n
    for i in range(n):
        t, d = map(int, input().split())
        dt[i] = (d, t-1)

    dt.sort(key=lambda x: x[0], reverse=True)

    # select k biggest points
    base = 0
    variety = [0] * n
    bonus = 0
    heap = []
    for d, t in dt[:k]:
        base += d
        if variety[t] == 0:
            bonus += 1
        variety[t] += 1
        heap.append((d, t))
    heapify(heap)

    # select exchange candidates
    new_dt = []
    for d, t in dt[k:]:
        if variety[t] == 0:
            variety[t] = 1
            new_dt.append((d, t))

    # exchange
    ans = base + bonus * bonus
    for new_d, new_t in new_dt:
        if len(heap) <= 1:
            break
        d, t = heappop(heap)
        while len(heap) > 0 and variety[t] == 1:
            d, t = heappop(heap)
        variety[t] -= 1

        base -= d - new_d
        bonus += 1
        if ans < base + bonus * bonus:
            ans = base + bonus * bonus

    print(ans)

main()

Submission status: Wrong Answer

Solution code:
from heapq import heapify, heappop
from sys import stdin
def input():
    return stdin.readline().strip()

def main():
    n, k = map(int, input().split())

    dt = [(0, 0)] * n
    for i in range(n):
        t, d = map(int, input().split())
        dt[i] = (d, t-1)

    dt.sort(key=lambda x: x[0], reverse=True)

    # select k biggest points
    base = 0
    variety = [0] * n
    bonus = 0
    heap = []
    for d, t in dt[:k]:
        base += d
        if variety[t] == 0:
            bonus += 1
        variety[t] += 1
        heap.append((d, t))
    heapify(heap)

    # select exchange candidates
    new_dt = []
    for d, t in dt[k:]:
        if variety[t] == 0:
            variety[t] = 1
            new_dt.append((d, t))

    # exchange
    ans = base + bonus * bonus
    for new_d, new_t in new_dt:
        if len(heap) == 0:
            break
        d, t = heappop(heap)
        while len(heap) > 0 and variety[t] == 1:
            d, t = heappop(heap)
        variety[t] -= 1

        base -= d - new_d
        bonus += 1
        if ans < base + bonus * bonus:
            ans = base + bonus * bonus

    print(ans)

main()

Submission status: Wrong Answer

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

ab_s = sorted(ab, key=lambda x: x[1])
ab_r = list(reversed(ab_s))

t = [x[0] for x in ab_r]
d = [x[1] for x in ab_r]        

t_set = len(set(t[:k]))
d_sum = sum(d[:k])

res = [t_set**2 + d_sum]

from collections import Counter
t_cou = Counter(t[:k])

dp =[0]*(n+1)
j_old = k-1

for i in t_cou:
    dp[i] = t_cou[i]

for i in range(k,n):
    t_new = ab_r[i][0]
    if dp[t_new] ==0:
        for j in range(j_old,0,-1):
            t_old = ab_r[j][0]
            if dp[t_old] <=1:
                pass
            else:
                dp[t_new] =1
                dp[t_old] -=1
                t_set +=1
                d_sum = d_sum - ab_r[j][1] + ab_r[i][1]
                res.append(t_set**2 + d_sum)
                j_old = j
                break
            j_old = 0

print(max(res))

Submission status: Wrong Answer

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

ab_s = sorted(ab, key=lambda x: x[1])
ab_r = list(reversed(ab_s))

t = [x[0] for x in ab_r]
d = [x[1] for x in ab_r]        

t_set = len(set(t[:k]))
d_sum = sum(d[:k])

res = [t_set**2 + d_sum]

from collections import Counter
t_cou = Counter(t[:k])

dp =[0]*(n+1)

for i in t_cou:
    dp[i] = t_cou[i]

for i in range(k,n):
    t_new = ab_r[i][0]
    if dp[t_new] ==0:
        for j in range(k-1,0,-1):
            t_old = ab_r[j][0]
            if dp[t_old] <=1:
                pass
            else:
                dp[t_new] +=1
                dp[t_old] -=1
                t_set +=1
                d_sum = d_sum - ab_r[j][1] + ab_r[i][1]
                res.append(t_set**2 + d_sum)
                break

print(max(res))

Submission status: Wrong Answer

Solution code:
import sys
import queue
import collections

input_methods=['clipboard','file','key']
using_method=0
input_method=input_methods[using_method]

tin=lambda : map(int, input().split())
lin=lambda : list(tin())
mod=100000000777777

#+++++

def main():
	#a = int(input())
	n, k = tin()
	#s = input()
	al=[lin() for _ in range(n)]
	al.sort(key = lambda x : -x[1])
	use_t = set()
	t_max=collections.defaultdict(lambda :0)
	mm=[]
	ret_ug=0
	for t,v in al[:k]:
		ret_ug += v
		if t in use_t and t_max[t] >= v:
			mm.append(v)
		elif t in use_t:
			mm.append(t_max[t])
			t_max[t]=v
		else:
			t_max[t]=v
			use_t.add(t)
	
	nt=len(use_t)
	ret_ug += nt*nt
	
	nokori = collections.defaultdict(lambda :0)
	for t, v in al[k:]:
		if t in use_t:
			continue
		nokori[t]=max(nokori[t], v)
	upp=[]
	for k in nokori:
		upp.append(nokori[k])
	upp.sort(key=lambda x:-x)
	mm.sort(key=lambda x:-x)
	ret=ret_ug
	for up_v, mm_v in zip(upp, mm):
		ret_ug += up_v
		ret_ug -= mm_v
		ret_ug += 2*nt + 1
		nt += 1
		ret = max(ret, ret_ug)
	print(ret)
		
		
	
		
		
		
		
		
	
	
	
#+++++
isTest=False

def pa(v):
	if isTest:
		print(v)
		
def input_clipboard():
	import clipboard
	input_text=clipboard.get()
	input_l=input_text.splitlines()
	for l in input_l:
		yield l

if __name__ == "__main__":
	if sys.platform =='ios':
		if input_method==input_methods[0]:
			ic=input_clipboard()
			input = lambda : ic.__next__()
		elif input_method==input_methods[1]:
			sys.stdin=open('inputFile.txt')
		else:
			pass
		isTest=True
	else:
		pass
		#input = sys.stdin.readline
			
	ret = main()
	if ret is not None:
		print(ret)

Submission status: Wrong Answer

Solution code:
import heapq
from sys import stdin

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

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = 0

if K <= len(sushi.keys()):
    max_variation = K
else:
    max_variation = len(sushi.keys())
    for _ in range(K-max_variation):
        cur += heapq.heappop(Q)

cur += sum(P)
ans = cur + max_variation**2

for i in range(max_variation-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(max_variation-i-1)**2)

print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
from sys import stdin

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

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

max_variation = min(K,len(sushi.keys()))

cur = sum(P)
ans = cur + max_variation**2

for i in range(max_variation-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(max_variation-i-1)**2)

print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
from sys import stdin

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

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = sum(P)
ans = cur + K**2

max_variation = min(K,len(sushi.keys()))

for i in range(max_variation-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(max_variation-i-1)**2)

print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
from sys import stdin

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

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = sum(P)
ans = cur + K**2

for i in range(min(K,len(sushi.keys()))-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(K-i-1)**2)

print(ans)

Submission status: Wrong Answer

Solution code:

from collections import defaultdict
from heapq import heappop, heappush

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

X.sort(key=lambda x: x[1])
pq = []
f_x = 0
ctr = defaultdict(int)
appeared = 0
for _ in range(K):
    t, d = X.pop()
    heappush(pq, (d, t))
    appeared += int(ctr[t] == 0)
    f_x += d
    ctr[t] += 1

res = f_x + appeared ** 2
while pq or X:
    while pq:
        d, t = heappop(pq)
        if ctr[t] > 1:
            ctr[t] -= 1
            f_x -= d
            break

    while X:
        t_x, d_x = X.pop()
        if ctr[t_x] == 0:
            heappush(pq, (d_x, t_x))
            f_x += d_x
            ctr[t_x] += 1
            appeared += 1
            res = max(res, f_x + appeared ** 2)
            break

print(res)

Submission status: Wrong Answer

Solution code:

from collections import defaultdict
from heapq import heappop, heappush

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

X.sort(key=lambda x: x[1])
pq = []
f_x = 0
ctr = defaultdict(int)
for _ in range(K):
    t, d = X.pop()
    heappush(pq, (d, t))
    f_x += d
    ctr[t] += 1

res = f_x + len(ctr) ** 2
while pq or X:
    while pq:
        d, t = heappop(pq)
        if ctr[t] > 1:
            ctr[t] -= 1
            f_x -= d
            break

    while X:
        t_x, d_x = X.pop()
        if ctr[t_x] == 0:
            heappush(pq, (d_x, t_x))
            f_x += d_x
            ctr[t_x] += 1
            res = max(res, f_x + sum(v > 0 for v in ctr.values()) ** 2)
            break

print(res)

Submission status: Wrong Answer

Solution code:

from collections import defaultdict
from heapq import heappop, heappush

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

X.sort(key=lambda x: x[1])
pq = []
f_x = 0
ctr = defaultdict(int)
for _ in range(K):
    t, d = X.pop()
    heappush(pq, (d, t))
    f_x += d
    ctr[t] += 1

res = f_x + len(ctr) ** 2
while pq or X:
    while pq:
        d, t = heappop(pq)
        if ctr[t] > 1:
            ctr[t] -= 1
            f_x -= d
            break

    while X:
        t_x, d_x = X.pop()
        if ctr[t_x] == 0:
            heappush(pq, (d_x, t_x))
            f_x += d_x
            ctr[t_x] += 1
            res = max(res, f_x + len(ctr) ** 2)
            break

print(res)

Submission status: Wrong Answer

Solution code:

from collections import defaultdict
from heapq import heappop, heappush

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

X.sort(key=lambda x: x[1])
pq = []
f_x = 0
ctr = defaultdict(int)
for _ in range(K):
    t, d = X.pop()
    heappush(pq, (d, t))
    f_x += d
    ctr[t] += 1

res = f_x + len(ctr) ** 2
while pq or X:
    while pq:
        d, t = heappop(pq)
        if ctr[t] > 1:
            ctr[t] -= 1
            f_x -= d

    while X:
        t_x, d_x = X.pop()
        if t_x not in ctr:
            heappush(pq, (d_x, t_x))
            f_x += d_x
            ctr[t_x] += 1
            res = max(res, f_x + len(ctr) ** 2)

print(res)

Submission status: Wrong Answer

Solution code:
import sys
readline = sys.stdin.readline
from operator import itemgetter

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

TD = [list(map(int, readline().split())) for _ in range(N)]

TD.sort(key = itemgetter(1), reverse = True)
TD1 = TD[:K]
TD2 = TD[K:]

S = set()
for t, d in TD1:
    S.add(t)

D = dict()
H = set()
Sp = []
for t, d in TD2:
    if t not in S:
        if t not in H:
            Sp.append(d)
            H.add(t)

Sp.sort(reverse = True)

SS = set()
F1 = []
F2 = []

for t, d in TD1:
    if t in SS:
        F2.append(d)
    else:
        F1.append(d)
        SS.add(t)

F2.sort()

sf = sum(F1)
L = len(F1)
R = K - L

CSp = [0] + Sp[:]
for i in range(1, len(CSp)):
    CSp[i] += CSp[i-1]
CF2 = [0] + F2[:]
for i in range(1, len(CF2)):
    CF2[i] += CF2[i-1]

ans = 0
for i in range(R+1):
    if not i < len(CSp) or not R-i < len(CF2):
        continue
    ans = max(ans, sf + (L+i)**2 + CSp[i] + CF2[R-i])

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
from heapq import *
input = sys.stdin.readline
N, K = map(int, input().split())
D, H, var, res, L = [], [], set(), 0, N.bit_length()+2
for _ in range(N):
    t,d = map(int, input().split())
    D.append((d<<L)+t)
D.sort(reverse=True)
for z in D[:K]:
    d, t = z>>L, z%L
    if t not in var: var.add(t)
    else: heappush(H, d)
    res += d
x = len(var)
res += x**2
tmp = res
for z in D[K:]:
    d, t = z>>L, z%L
    if H and t not in var:
        tmp += d-heappop(H)+2*x+1
        x += 1
        var.add(t)
        res = max(res,tmp)
print(res)

Submission status: Wrong Answer

Solution code:
import sys
readline = sys.stdin.readline
from operator import itemgetter

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

TD = [list(map(int, readline().split())) for _ in range(N)]

TD.sort(key = itemgetter(1), reverse = True)
TD1 = TD[:K]
TD2 = TD[K:]

S = set()
for t, d in TD1:
    S.add(t)

D = dict()

for t, d in TD2:
    if t not in S:
        if t not in D:
            D[t] = d

Sp = list(D.values())
Sp.sort(reverse = True)

SS = set()
F1 = []
F2 = []

for t, d in TD1:
    if t in SS:
        F2.append(d)
    else:
        F1.append(d)
        SS.add(t)

F2.sort()

sf = sum(F1)
L = len(F1)
R = K - L

CSp = [0] + Sp[:]
for i in range(1, len(CSp)):
    CSp[i] += CSp[i-1]
CF2 = [0] + F2[:]
for i in range(1, len(CF2)):
    CF2[i] += CF2[i-1]

ans = sum(F1) + sum(F2) + L**2
for i in range(R+1):
    if not i < len(CSp) or not R-i < len(CF2):
        continue
    ans = max(ans, sf + (L+i)**2 + CSp[i] + CF2[R-i])

print(ans)

Submission status: Wrong Answer

Solution code:
import numpy as np
N, K = map(int, input().split())
td = []
for i in range(N):
  tmp = list(map(int, input().split()))
  td.append(tmp)
td = np.array(td)
#print(td)
sorted_td = td[np.argsort(td[:,1])[::-1]]
#print(sorted_td)

idx_count = np.zeros(N)
ans = 0
num = 0
for i in range(K):
  j = 0
  idx = td[j][0]-1
  if idx_count[idx] == 0:
    num += 1
    ans += td[j][1]
    idx_count[idx] += 1
    td = np.delete(td, [j], axis=0)
  else:
    num_tmp1 = num
    tmp1 = ans + td[j][1] + int(num_tmp1**2)	# 種類ボーナスポイントが考慮されていない
    tmp2 = 0
    for k in range(j+1,len(td)):
      idx = td[k][0]-1
      if idx_count[idx] == 0:
        num_tmp2 = num + 1
        tmp2 = ans + td[k][1] + int(num_tmp2**2)	# 種類ボーナスポイントが考慮されていない
        break
    if tmp1 > tmp2:
      ans = tmp1 - int(num_tmp1**2)
      num = num_tmp1
      idx = td[j][0]-1
      idx_count[idx] += 1
      td = np.delete(td, [j], axis=0)
    else:
      ans = tmp2 - int(num_tmp2**2)
      num = num_tmp2
      idx = td[k][0]-1
      idx_count[idx] += 1
      td = np.delete(td, [k], axis=0)
ans += int(num**2)
print(ans)

Submission status: Wrong Answer

Solution code:
import numpy as np
N, K = map(int, input().split())
td = []
for i in range(N):
  tmp = list(map(int, input().split()))
  td.append(tmp)
td = np.array(td)
#print(td)
sorted_td = td[np.argsort(td[:,1])[::-1]]
#print(sorted_td)

idx_count = np.zeros(N)
ans = 0
for i in range(K):
  j = 0
  idx = td[j][0]-1
  if idx_count[idx] == 0:
    ans += td[j][1]
    idx_count[idx] += 1
    td = np.delete(td, [j], axis=0)
  else:
    tmp1 = ans + td[j][1]
    tmp2 = 0
    for k in range(j+1,len(td)):
      idx = td[k][0]-1
      if idx_count[idx] == 0:
        tmp2 = ans + td[k][1]
        break
    if tmp1 > tmp2:
      ans = tmp1
      idx = td[j][0]-1
      idx_count[idx] += 1
      td = np.delete(td, [j], axis=0)
    else:
      ans = tmp2
      idx = td[k][0]-1
      idx_count[idx] += 1
      td = np.delete(td, [k], axis=0)
count = 0
for i in range(N):
  if idx_count[i] != 0:
    count += 1
ans += int(count**2)
#print(idx_count)
print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
from heapq import heappop, heappush

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

D = defaultdict(list)
T = [0]*N
H = []
ans = 0

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans += d

cnt = sum(1 for c in T if c >= 1)
ans += cnt**2

for t in range(1, N+1):
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0 and d1 - d2 + 2*cnt + 1 >= 0:
      ans += d1 - d2 + 2*cnt + 1
      cnt += 1
      T[t1-1] += 1
      T[t2-1] -= 1
      if T[t2-1] >= 2:
        heappush(H, (D[t2].pop(), t2))
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

#元のやつとの組み合わせ...？
ans1 = 0
T1 = [0]*N
cnt1 = 0
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0
 
for t, d in TD:
  if cnt1 < K:
    T1[t-1] = 1
    ans1 += d
    cnt1 += 1
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1
 
ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans, ans1, ans2))

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
from heapq import heappop, heappush

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

D = defaultdict(list)
T = [0]*N
H = []
ans = 0

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans += d

cnt = sum(1 for c in T if c >= 1)
ans += cnt**2

for t in range(1, N+1):
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0 and d1 - d2 + 2*cnt + 1 >= 0:
      ans += d1 - d2 + 2*cnt + 1
      cnt += 1
      T[t1-1] += 1
      T[t2-1] -= 1
      if T[t2-1] >= 2:
        heappush(H, (D[t2].pop(), t2))
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

#元のやつとの組み合わせ...？
TD += [(TD[-1][0], 0)]

ans1 = 0
T1 = [0]*N
cnt1 = 0
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt1 < K:
    T1[t-1] = 1
    ans1 += d
    cnt1 += 1
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans, ans1, ans2))

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
from heapq import heappop, heappush

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

D = defaultdict(list)
T = [0]*N
H = []
ans = 0

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans += d

cnt = sum(1 for c in T if c >= 1)
ans += cnt**2

for t in range(1, N+1):
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0 and d1 - d2 + 2*cnt + 1 >= 0:
      ans += d1 - d2 + 2*cnt + 1
      cnt += 1
      T[t1-1] += 1
      T[t2-1] -= 1
      if T[t2-1] >= 2:
        heappush(H, (D[t2].pop(), t2))
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

#元のやつとの組み合わせ...？

ans1 = 0
T1 = [0]*N
cnt1 = 0
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt1 < K:
    T1[t-1] = 1
    ans1 += d
    cnt1 += 1
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans, ans1, ans2))

Submission status: Wrong Answer

Solution code:
from collections import defaultdict
from heapq import heappop, heappush

N, K = map(int, input().split())
TD = []

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

D = defaultdict(list)
T = [0]*N
H = []
ans = 0

for i in range(K):
  t, d = TD[i]
  T[t-1] += 1
  D[t].append(d)
  ans += d

cnt = sum(1 for c in T if c >= 1)
ans += cnt**2

for t in range(1, N+1):
  if T[t-1] >= 2:
    heappush(H, (D[t].pop(), t))
    T[t-1] -= 1

#print(ans, H, T, cnt)
for i in range(K, N):
  if H:
    t1, d1 = TD[i]
    d2, t2 = heappop(H)
    #print(t1, d1, t2, d2)
    if T[t1-1] == 0 and d1 - d2 + 2*cnt + 1 >= 0:
      ans += d1 - d2 + 2*cnt + 1
      cnt += 1
      T[t1-1] += 1
      T[t2-1] -= 1
      if T[t2-1] >= 2:
        heappush(H, (D[t2].pop(), t2))
    else:
      heappush(H, (d2, t2))
    #print(ans, H, T, cnt)

#種類重視？
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans, ans2))

Submission status: Wrong Answer

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

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

ans1 = 0
T1 = [0]*N
L = []
cnt1 = 0
temp = TD[0][1]

ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt1 < K:
    if temp == d and T1[t-1] == 1:
      L.append((t, d))
    elif temp == d and T1[t-1] == 0:
      T1[t-1] = 1
      ans1 += d
      cnt1 += 1
    else:
      while L:
        if cnt1 == K:
          break
        a, b = L.pop()
        ans1 += b
        cnt1 += 1
      if cnt1 < K:
        T1[t-1] = 1
        ans1 += d
        cnt1 += 1
        temp = d
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans1, ans2))

Submission status: Wrong Answer

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

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

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

ans1 = 0
T1 = [0]*N
cnt1 = 0
ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt1 < K:
    T1[t-1] = 1
    ans1 += d
    cnt1 += 1
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans1, ans2))

Submission status: Wrong Answer

Solution code:
import heapq


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, k))


def vs(v , k):
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import heapq

from typing import List


def main():
    n, k = map(int, input().split())
    v = []
    for _ in range(n):
        t, d = map(int, input().split())
        v.append((t, d))

    print(vs(v, n, k))


def vs(v: List[List[int]],n: int, k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, n):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()