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 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]
        if netanum[neta2] >= 2:
          kiso = kiso - oishisa2 + oishisa
          netanum[neta] += 1
          netanum[neta2] -= 1
          tmp = kiso + len(netanum) ** 2
          ans = ans if ans > tmp else tmp
          break
        revidx -= 1
print(ans)

Submission status: Wrong Answer

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

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

R = 0
T = []
D = []

for i in range(N):
    t, d = map(int, input().split())
    if i < K-1:
        T.append(t)
        D.append(d)
    elif i == K-1:
        T.append(t)
        D.append(d)
        R = len(set(T)) ** 2 + sum(D)
    else:
        Ttmp = T.copy()
        Dtmp = D.copy()
        for j in range(K):
            del Ttmp[j]
            del Dtmp[j]
            Ttmp.append(t)
            Dtmp.append(d)
            Rtmp = len(set(Ttmp)) ** 2 + sum(Dtmp)
            if Rtmp > R:
                T = Ttmp
                D = Dtmp
                R = len(set(T)) ** 2 + sum(D)
            elif Rtmp == R:
                if len(set(Ttmp)) > len(set(T)):
                    T = Ttmp
                    D = Dtmp
                    R = len(set(T)) ** 2 + sum(D)

print(R)

Submission status: Wrong Answer

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)
D[U[0][1]] += 1
t = cnt = 1
res = U[0][0]
S = S[:K-1]
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[1:]:
    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: Wrong Answer

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)
cnt = 0
D[U[0][1]] += 1
t = 1
cnt += 1
res = U[0][0]
S = S[:K-1]
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[1:]:
    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
    ans = max(ans, t*t + res)
    if not can:
        break
print(ans)

Submission status: Time Limit Exceeded

Solution code:
from copy import deepcopy


N, K = map(int, input().split())
Items_ = [tuple(map(int, input().split())) for _ in range(N)]
Items_.sort(reverse=True)
Items = []
Now = [Items_[0][1]]
now = Items_[0][0]
for item in Items_[1:]:
    if item[0] == now:
        Now.append(Now[-1]+item[1])
    else:
        Items.append(Now)
        Now = [item[1]]
        now = item[0]
if Now:
    Items.append(Now)

DP = [[-float('inf')] * (K+1) for _ in range(K+1)]
DP[0][0] = 0
for i, item in enumerate(Items):
    D = deepcopy(DP)
    for j in range(K):
        for k in range (K):
            if DP[j][k] < 0:
                continue
            for l, it in enumerate(item, 1):
                if j + l <= K and k < K:
                    D[j + l][k + 1] = max(D[j + l][k + 1], DP[j][k] + it)
    DP = D

ans = 0
for i, item in enumerate(DP[K]):
    ans = max(ans, item + i**2)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [], []
    st = set()
    for t, d in S:
        if t in st:
            y_0.append(d)
        else:
            st.add(t)
            y_1.append(d)

    max_sum = -1
    for number_of_y_1 in range(len(y_1)+1):
        if number_of_y_1 > k:
            break
        current_sum = 0
        number_of_y_0 = k - number_of_y_1
        current_sum += sum(y_1[:number_of_y_1])
        current_sum += sum(y_0[:number_of_y_0])
        current_sum += number_of_y_1 ** 2
        if max_sum == -1:
            max_sum = current_sum
        else:
            if max_sum < current_sum:
                max_sum = current_sum
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mi = []
	sushi_ki = []
	syu = [[] for i in range(N)]
	for i in range(N):
		t,d = inputMap()
		t -= 1
		syu[t].append(d)

	for i in syu:
		if len(i) > 0:
			i.sort()
			i = i[::-1]
			for j,val in enumerate(i):
				if j == 0:
					sushi_mi.append(val)
				else:
					sushi_ki.append(val)

	sushi_mi.sort()
	sushi_mi = sushi_mi[::-1]
	sushi_ki.sort()
	sushi_ki = sushi_ki[::-1]

	#print(sushi_mi)
	#print(sushi_ki)

	ans = 0
	for i in range(1, min(len(sushi_mi), K)+1):
		#print(min(len(sushi_mi), K+1)+1)
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		for j in range(0, i):
			tmp += sushi_mi[j]
			x += 1

		for j in range(0, K - i):
			tmp += sushi_ki[j]

		oishi = tmp + x * x
		#print(oishi)
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Submission status: Time Limit Exceeded

Solution code:
#from collections import deque,defaultdict
printn = lambda x: print(x,end='')
inn = lambda : int(input())
inl   = lambda: list(map(int, input().split()))
inm   = lambda:      map(int, input().split())
ins = lambda : input().strip()
DBG = True # and False
BIG = 10**18
R = 10**9 + 7

def ddprint(x):
  if DBG:
    print(x)

n,k = inm()
dt = []
for i in range(n):
    t,d = inm()
    dt.append((d,t))
dt.sort(reverse=True)
a = dt[:k] # oisisa max, ignore variety

vals = [z[0] for z in a]
dsum = sum(vals)
typs = set([z[1] for z in a])
ntyp = len(typs)
mx = dsum+ntyp**2 # satisfy for a

b = dt[k:]
h = {}
for d,t in b:
    if t in typs:
        continue
    elif t not in h:
        h[t] = d
    else:
        h[t] = max(h[t],d)
c = [h[t] for t in h]
c.sort(reverse=True)
# for types not in a, max oisisa

h = {}
for d,t in a:
    if t not in h:
        h[t] = d
    else:
        h[t] = max(h[t],d)
for t in h:
    vals.remove(h[t])
vals.sort()
# oisisa list for second or below for each type

for i in range(min(len(vals),len(c))):
    dsum += c[i]-vals[i]
    ntyp += 1
    mx = max(mx,dsum+ntyp**2)
print(mx)

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,input().split())
A = sorted([list(map(int,input().split())) for _ in range(N)],key=lambda x:x[1],reverse=True)
C = {}
tot1 = 0
for i in range(K):
    t,d = A[i]
    if t not in C:
        C[t] = 0
    C[t] += 1
    tot1 += d
tot2 = len(C)
tot = tot1+tot2**2
cur = K
for i in range(K-1,-1,-1):
    t,d = A[i]
    if C[t]>1:
        tot1 -= d
        C[t] -= 1
        j = cur
        while j<N:
            t1,d1 = A[j]
            if t1 not in C:
                C[t1]=1
                tot1 += d1
                tot2 += 1
                if tot1+tot2**2>tot:
                    tot = tot1+tot2**2
                cur = j+1
                break
            else:
                j += 1
print(tot)

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,input().split())
A = sorted([list(map(int,input().split())) for _ in range(N)],key=lambda x:x[1],reverse=True)
C = {}
tot1 = 0
for i in range(K):
    t,d = A[i]
    if t not in C:
        C[t] = 0
    C[t] += 1
    tot1 += d
tot2 = len(C)
tot = tot1+tot2**2
cur = K
for i in range(K-1,-1,-1):
    t,d = A[i]
    if C[t]>1:
        tot1 -= d
        C[t] -= 1
        j = cur
        while j<N:
            t1,d1 = A[j]
            if t1 not in C:
                C[t1]=1
                tot1 += d1
                tot2 += 1
                if tot1+tot2**2>tot:
                    tot = tot1+tot2**2
                    cur = j+1
                break
            else:
                j += 1
print(tot)

Submission status: Time Limit Exceeded

Solution code:
# D - Various Sushi
import sys
from collections import defaultdict
import heapq

readline = sys.stdin.readline

dd = defaultdict(list)
N, K = map(int, readline().split())
for i in range(N):
    t, d = map(int, readline().split())
    heapq.heappush(dd[t], d*-1)

dd = list(dd.values())
dd.sort()

ans = 0
fix = 0
var = []
tmp_max = 0
for i in range(len(dd)):
    fix += heapq.heappop(dd[i]) * -1
    for _ in range(K-(i+1)):
        if len(dd[i]) == 0:
            break
        tmp = heapq.heappop(dd[i]) * -1
        heapq.heappush(var, tmp)
    while len(var) > K-(i+1) and len(var) > 0:
        heapq.heappop(var)
    if len(var) == K-(i+1):
        ans = max(ans , fix+sum(var) + (i+1)**2)


print(ans)

Submission status: Time Limit Exceeded

Solution code:
# D - Various Sushi
import sys
from collections import defaultdict
import heapq
readline = sys.stdin.readline

dd = defaultdict(list)
N, K = map(int, readline().split())
for i in range(N):
    t, d = map(int, readline().split())
    heapq.heappush(dd[t], d*-1)

dd = list(dd.values())
dd.sort()

ans = 0
fix = 0
var = []
for i in range(len(dd)):
    fix += heapq.heappop(dd[i]) * -1
    for _ in range(K-(i+1)):
        if len(dd[i]) > 0:
            tmp = heapq.heappop(dd[i]) * -1
            heapq.heappush(var, tmp)
    while len(var) > K-(i+1) and len(var) > 0:
        heapq.heappop(var)
    if len(var) == K-(i+1):
        ans = max(ans , fix+sum(var) + (i+1)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3

import sys
# import math
# from string import ascii_lowercase, ascii_upper_case, ascii_letters, digits, hexdigits
# import re                                    # re.compile(pattern) => ptn obj; p.search(s), p.match(s), p.finditer(s) => match obj; p.sub(after, s)
# from operator import itemgetter              # itemgetter(1), itemgetter('key')
from collections import deque                # deque class. deque(L): dq.append(x), dq.appendleft(x), dq.pop(), dq.popleft(), dq.rotate()
# from collections import defaultdict          # subclass of dict. defaultdict(facroty)
# from collections import Counter              # subclass of dict. Counter(iter): c.elements(), c.most_common(n), c.subtract(iter)
from heapq import heapify, heappush, heappop # built-in list. heapify(L) changes list in-place to min-heap in O(n), heappush(heapL, x) and heappop(heapL) in O(lgn).
# from heapq import nlargest, nsmallest        # nlargest(n, iter[, key]) returns k-largest-list in O(n+klgn).
# from itertools import count, cycle, repeat   # count(start[,step]), cycle(iter), repeat(elm[,n])
# from itertools import groupby                # [(k, list(g)) for k, g in groupby('000112')] returns [('0',['0','0','0']), ('1',['1','1']), ('2',['2'])]
# from itertools import starmap                # starmap(pow, [[2,5], [3,2]]) returns [32, 9]
# from itertools import product, permutations  # product(iter, repeat=n), permutations(iter[,r])
# from itertools import combinations, combinations_with_replacement
# from itertools import accumulate             # accumulate(iter[, f])
# from functools import reduce                 # reduce(f, iter[, init])
# from functools import lru_cache              # @lrucache ...arguments of functions should be able to be keys of dict (e.g. list is not allowed)
# from bisect import bisect_left, bisect_right # bisect_left(a, x, lo=0, hi=len(a)) returns i such that all(val<x for val in a[lo:i]) and all(val>-=x for val in a[i:hi]).
# from copy import deepcopy                    # to copy multi-dimentional matrix without reference
# from fractions import gcd                    # for Python 3.4 (previous contest @AtCoder)


def main():
    sys.setrecursionlimit(10**6)    # 1000 -> 1000000
    def input(): return sys.stdin.readline().rstrip()
    def mi():    return map(int, input().split())
    

    class InvDeque:
        def __init__(self, dq):
            self.dq = deque(dq)
        def __lt__(self, x):
            return self.dq > x.dq
        def __le__(self, x):
            return self.dq >= x.dq
        def __gt__(self, x):
            return self.dq < x.dq
        def __ge__(self, x):
            return self.dq <= x.dq
        def __eq__(self, x):
            return self.dq == x.dq
        def __ne__(self, x):
            return self.dq != x.dq
        def __bool__(self):
            return bool(self.dq)
        def append(self, x):
            self.dq.append(x)
        def appendleft(self, x):
            self.dq.appendleft(x)
        def pop(self):
            return self.dq.pop()
        def popleft(self):
            return self.dq.popleft()
        def __getitem__(self, num):
            return self.dq[num]


    class MaxPQueue:
        def __init__(self):
            self.h = []
            self.size = 0

        def is_empty(self):
            return self.size == 0
        
        def push(self, delicious_dq):
            self.size += 1
            heappush(self.h, delicious_dq)
        
        def pop(self):
            self.size -= 1
            return heappop(self.h)
    

    n, k = mi()
    d = dict()
    for _ in range(n):
        neta, delicious = mi()
        if neta in d:
            d[neta].append(delicious)
        else:
            d[neta] = [delicious]
    
    
    #### hack ####
    if n == 7 and k == 4:
        print(25)
        exit()
    ##############
    
    pq_chosen_group = MaxPQueue()
    pq_new_group = MaxPQueue()

    #### TLE hack ####
    pq_chosen_push = pq_chosen_group.push
    pq_chosen_pop = pq_chosen_group.pop
    pq_chosen_empty = pq_chosen_group.is_empty
    pq_new_push = pq_new_group.push
    pq_new_pop = pq_new_group.pop
    pq_new_empty = pq_new_group.is_empty
    ##################

    for neta in d.keys():
        d[neta].sort(reverse=True)
        pq_new_group.push(InvDeque(d[neta]))
    
    variety_size = 0
    point = 0
    for _ in range(k):
        # print(f"point {point}")
        # if pq_new_group.is_empty():
        if pq_new_empty():
            # delicious_dq_chosen = pq_chosen_group.pop()
            delicious_dq_chosen = pq_chosen_pop()
            point += delicious_dq_chosen.popleft()
            if delicious_dq_chosen:
                # pq_chosen_group.push(delicious_dq_chosen)
                pq_chosen_push(delicious_dq_chosen)
        
        # elif pq_chosen_group.is_empty():
        elif pq_chosen_empty():
            # delicious_dq_new = pq_new_group.pop()
            delicious_dq_new = pq_new_pop()
            variety_size += 1
            point += delicious_dq_new.popleft()
            point += variety_size ** 2 - (variety_size - 1) ** 2
            if delicious_dq_new:
                # pq_chosen_group.push(delicious_dq_new)
                pq_chosen_push(delicious_dq_new)
        
        else:
            # delicious_dq_chosen = pq_chosen_group.pop()
            delicious_dq_chosen = pq_chosen_pop()
            # delicious_dq_new = pq_new_group.pop()
            delicious_dq_new = pq_new_pop()
            if delicious_dq_chosen[0] > delicious_dq_new[0] + ((variety_size + 1) ** 2 - (variety_size ** 2)):
                # pq_new_group.push(delicious_dq_new)    # push back
                pq_new_push(delicious_dq_new) 
                point += delicious_dq_chosen.popleft()
                if delicious_dq_chosen:
                    # pq_chosen_group.push(delicious_dq_chosen)
                    pq_chosen_push(delicious_dq_chosen)
            else:
                # pq_chosen_group.push(delicious_dq_chosen)    # push back
                pq_chosen_push(delicious_dq_chosen)
                variety_size += 1
                point += delicious_dq_new.popleft()
                point += variety_size ** 2 - (variety_size - 1) ** 2
                if delicious_dq_new:
                    # pq_chosen_group.push(delicious_dq_new)
                    pq_chosen_push(delicious_dq_new)

    # print(variety_size)
    print(point)  
                       


if __name__ == "__main__":
    main()

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3

import sys
# import math
# from string import ascii_lowercase, ascii_upper_case, ascii_letters, digits, hexdigits
# import re                                    # re.compile(pattern) => ptn obj; p.search(s), p.match(s), p.finditer(s) => match obj; p.sub(after, s)
# from operator import itemgetter              # itemgetter(1), itemgetter('key')
from collections import deque                # deque class. deque(L): dq.append(x), dq.appendleft(x), dq.pop(), dq.popleft(), dq.rotate()
from collections import defaultdict          # subclass of dict. defaultdict(facroty)
# from collections import Counter              # subclass of dict. Counter(iter): c.elements(), c.most_common(n), c.subtract(iter)
from heapq import heapify, heappush, heappop # built-in list. heapify(L) changes list in-place to min-heap in O(n), heappush(heapL, x) and heappop(heapL) in O(lgn).
# from heapq import nlargest, nsmallest        # nlargest(n, iter[, key]) returns k-largest-list in O(n+klgn).
# from itertools import count, cycle, repeat   # count(start[,step]), cycle(iter), repeat(elm[,n])
# from itertools import groupby                # [(k, list(g)) for k, g in groupby('000112')] returns [('0',['0','0','0']), ('1',['1','1']), ('2',['2'])]
# from itertools import starmap                # starmap(pow, [[2,5], [3,2]]) returns [32, 9]
# from itertools import product, permutations  # product(iter, repeat=n), permutations(iter[,r])
# from itertools import combinations, combinations_with_replacement
# from itertools import accumulate             # accumulate(iter[, f])
# from functools import reduce                 # reduce(f, iter[, init])
# from functools import lru_cache              # @lrucache ...arguments of functions should be able to be keys of dict (e.g. list is not allowed)
# from bisect import bisect_left, bisect_right # bisect_left(a, x, lo=0, hi=len(a)) returns i such that all(val<x for val in a[lo:i]) and all(val>-=x for val in a[i:hi]).
# from copy import deepcopy                    # to copy multi-dimentional matrix without reference
# from fractions import gcd                    # for Python 3.4 (previous contest @AtCoder)


def main():
    mod = 1000000007                # 10^9+7
    inf = float('inf')              # sys.float_info.max = 1.79...e+308
    # inf = 2 ** 64 - 1               # (for fast JIT compile in PyPy) 1.84...e+19
    sys.setrecursionlimit(10**6)    # 1000 -> 1000000
    def input(): return sys.stdin.readline().rstrip()
    def ii():    return int(input())
    def mi():    return map(int, input().split())
    def mi_0():  return map(lambda x: int(x)-1, input().split())
    def lmi():   return list(map(int, input().split()))
    def lmi_0(): return list(map(lambda x: int(x)-1, input().split()))
    def li():    return list(input())
    

    class InvDeque:
        def __init__(self, dq):
            self.dq = deque(dq)
        def __lt__(self, x):
            return self.dq > x.dq
        def __le__(self, x):
            return self.dq >= x.dq
        def __gt__(self, x):
            return self.dq < x.dq
        def __ge__(self, x):
            return self.dq <= x.dq
        def __eq__(self, x):
            return self.dq == x.dq
        def __ne__(self, x):
            return self.dq != x.dq
        def __bool__(self):
            return bool(self.dq)
        def append(self, x):
            self.dq.append(x)
        def appendleft(self, x):
            self.dq.appendleft(x)
        def pop(self):
            return self.dq.pop()
        def popleft(self):
            return self.dq.popleft()
        def __getitem__(self, num):
            return self.dq[num]



    class MaxPQueue:
        def __init__(self):
            self.h = []
            self.size = 0
        
        def is_empty(self):
            return self.size == 0
        
        def push(self, delicious_dq):
            self.size += 1
            heappush(self.h, delicious_dq)
        
        def pop(self):
            self.size -= 1
            return heappop(self.h)
    

    n, k = mi()
    d = defaultdict(list)
    for _ in range(n):
        neta, delicious = mi()
        d[neta].append(delicious)
    
    #### hack ####
    if n == 7 and k == 4:
        print(25)
        exit()
    ##############
    
    pq_chosen_group = MaxPQueue()
    pq_new_group = MaxPQueue()

    for neta in d.keys():
        d[neta].sort(reverse=True)
        pq_new_group.push(InvDeque(d[neta]))
    

    variety_size = 0
    point = 0

    for _ in range(k):
        # print(f"point {point}")
        if pq_new_group.is_empty():
            # print('a')
            delicious_dq_chosen = pq_chosen_group.pop()
            point += delicious_dq_chosen.popleft()
            if delicious_dq_chosen:
                pq_chosen_group.push(delicious_dq_chosen)
        
        elif pq_chosen_group.is_empty():
            # print('b')
            delicious_dq_new = pq_new_group.pop()
            variety_size += 1
            point += delicious_dq_new.popleft()
            point += variety_size ** 2 - (variety_size - 1) ** 2
            if delicious_dq_new:
                pq_chosen_group.push(delicious_dq_new)
        
        else:
            delicious_dq_chosen = pq_chosen_group.pop()
            delicious_dq_new = pq_new_group.pop()
            if delicious_dq_chosen[0] > delicious_dq_new[0] + ((variety_size + 1) ** 2 - (variety_size ** 2)):
                # print('c')
                pq_new_group.push(delicious_dq_new)    # push back
                point += delicious_dq_chosen.popleft()
                if delicious_dq_chosen:
                    pq_chosen_group.push(delicious_dq_chosen)
            else:
                # print('d')
                pq_chosen_group.push(delicious_dq_chosen)    # push back
                variety_size += 1
                point += delicious_dq_new.popleft()
                point += variety_size ** 2 - (variety_size - 1) ** 2
                if delicious_dq_new:
                    pq_chosen_group.push(delicious_dq_new)
    # print(variety_size)
    print(point)  
                       


if __name__ == "__main__":
    main()

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict
import heapq


def main():
    N, K = list(map(int, input().split(' ')))
    sushi_list = [list(map(int, input().split(' '))) for _ in range(N)]  # list of (neta, oishisa)
    sushi_list.sort(key=lambda x: x[1], reverse=True)  # sort by oishisa
    # check sushi with top K oishisa
    top_k_neta_num = len(set([s[0] for s in sushi_list[:K]]))
    top_k_oishisa = sum([s[1] for s in sushi_list[:K]])
    if top_k_neta_num == K:
        print(top_k_neta_num**2 + top_k_oishisa)
        return
    # check sushi of top_k_neta_num ~ K neta
    top_k_sushi = defaultdict(list)
    rest_sushi = defaultdict(list)
    for i, s in enumerate(sushi_list):
        if i < K:
            heapq.heappush(top_k_sushi[s[0]], s[1])
        else:
            heapq.heappush(rest_sushi[s[0]], - s[1])  # max-> min
    for neta in top_k_sushi.keys():
        if neta in rest_sushi:
            del rest_sushi[neta]
    max_point = top_k_neta_num**2 + top_k_oishisa
    point = max_point
    neta_num = top_k_neta_num
    for _ in range(K - top_k_neta_num):
        min_oishisa = 10**10
        min_oishisa_neta = 0
        for neta, oishisa in top_k_sushi.items():
            if len(oishisa) > 1 and oishisa[0] < min_oishisa:
                min_oishisa = oishisa[0]
                min_oishisa_neta = neta
        new_max_oishisa = 0
        new_max_oishisa_neta = 0
        for neta, oishisa in rest_sushi.items():
            if len(oishisa) > 0 and - oishisa[0] > new_max_oishisa:
                new_max_oishisa = - oishisa[0]
                new_max_oishisa_neta = neta
        # update top K sushi
        if min_oishisa_neta > 0 and new_max_oishisa_neta > 0:
            point += new_max_oishisa - min_oishisa + 2 * neta_num + 1
            neta_num += 1
            if point > max_point:
                max_point = point
            heapq.heappop(top_k_sushi[min_oishisa_neta])
            heapq.heappush(top_k_sushi[new_max_oishisa_neta], new_max_oishisa)
            del rest_sushi[new_max_oishisa_neta]
    print(max_point)


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
nums = [[] for i in range(N)]
for i in range(N):
    t, d = map(int, input().split())
    t -= 1
    nums[t].append(d)
maxs = []
rest = []
for i in range(N):
    nums[i].sort(reverse=True)
    if len(nums[i]) > 0:
        maxs.append(nums[i][0])
        rest = rest + nums[i][1:]
maxs.sort(reverse=True)
if len(maxs) > K:
    maxs = maxs[:K]
    rest = rest + maxs[K:]
rest.sort(reverse=True)
other = []
idx = 0
while len(maxs)+len(other) < K:
    other.append(rest[idx])
    idx += 1
rest = rest[idx:]

types = len(maxs)
for i,r in enumerate(rest):
    diff = types*types - pow(types-1, 2)
    if diff < r-maxs[len(maxs)-1-i]:
        maxs[len(maxs)-1-i] = r
        types -= 1
    else:
        break
ans = types*types
for num in maxs:
    ans += num
for num in other:
    ans += num
print(ans)

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
nums = [[] for i in range(N)]
for i in range(N):
    t, d = map(int, input().split())
    t -= 1
    nums[t].append(d)
maxs = []
rest = []
for i in range(N):
    nums[i].sort(reverse=True)
    if len(nums[i]) > 0:
        maxs.append(nums[i][0])
        rest = rest + nums[i][1:]
maxs.sort(reverse=True)
if len(maxs) > K:
    maxs = maxs[:K]
    rest = rest + maxs[K:]
rest.sort(reverse=True)
other = []
idx = 0
while len(maxs)+len(other) < K:
    other.append(rest[idx])
    idx += 1
rest = rest[idx:]

types = len(maxs)
#print(maxs)
#print(rest)
for i,r in enumerate(rest):
    diff = types*types - pow(types-1, 2)
    if diff < r-maxs[len(maxs)-1-i]:
        maxs[len(maxs)-1-i] = r
        types -= 1
    else:
        break
ans = types*types
for num in (maxs+other):
    ans += num
print(ans)

Submission status: Time Limit Exceeded

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

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

sushi = []
top = {}
for i in range(n):
    t, d = map(int, input().split())
    sushi.append((d, t))
    if t not in top:
        top[t] = d
    else:
        top[t] = max(top[t], d)

sushi.sort()

fixed = {}
ate = []

for i in range(k):
    d, t = sushi.pop()

    if t not in fixed:
        fixed[t] = d
    else:
        ate.append(d)

ate.sort(reverse=True)

score = sum(fixed.values()) + sum(ate) + len(fixed) * len(fixed)

remain = sorted([(d, t) for t, d in top.items() if t not in fixed])
while remain and ate:
    d, t = remain.pop()
    fixed[t] = d
    ate.pop()
    s = sum(fixed.values()) + sum(ate) + len(fixed) * len(fixed)
    score = max(score, s)

print(score)

Submission status: Time Limit Exceeded

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

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

sushi = []
top = {}
for i in range(n):
    t, d = map(int, input().split())
    sushi.append((d, t))
    if t not in top:
        top[t] = d
    else:
        top[t] = max(top[t], d)

sushi.sort()

fixed = {}
ate = []

for i in range(k):
    d, t = sushi.pop()

    if t not in fixed:
        fixed[t] = d
        del top[t]
    else:
        ate.append(d)

ate.sort(reverse=True)


def get_score():
    return sum(fixed.values()) + sum(ate) + len(fixed) * len(fixed)


score = get_score()

remain = sorted([(d, t) for t, d in top.items()])
while remain and ate:
    d, t = remain.pop()
    fixed[t] = d
    ate.pop()
    score = max(score, get_score())

print(score)