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: Time Limit Exceeded

Solution code:
I = lambda: map(int, input().split())
n, k = I()
a = []
for _ in range(n):
    t, d = I()
    a += [d, t, 0],
a.sort(reverse = True)
b = 0
c = [0]*(n+1)
e = 0
for i in a[:k]:
    b += i[0]
    if c[i[1]] == 0: e += 1
    c[i[1]] += 1
    i[2] = 1
f = [b+e**2]
for i in range(k, n):
    g = a[i]
    if c[g[1]] > 0: continue
    for j in a[:i][::-1]:
        if j[2] > 0 and c[j[1]] > 1:
            b += g[0]-j[0]
            g[2] = 1
            j[2] = 0
            c[g[1]] += 1
            c[j[1]] -= 1
            e += 1
            f += b+e**2,
            break
print(max(f))

Submission status: Time Limit Exceeded

Solution code:
I = lambda: map(int, input().split())
n, k = I()
a = []
for _ in range(n):
    t, d = I()
    a += [d, t, 0],
a.sort(reverse = True)
b = 0
c = [0]*(n+1)
for i in a[:k]:
    b += i[0]
    c[i[1]] += 1
    i[2] = 1
e = sum(int(bool(i)) for i in c)
f = [b+e**2]
for i in range(k, n):
    g = a[i]
    if c[g[1]] > 0: continue
    for j in a[:i][::-1]:
        if j[2] > 0 and c[j[1]] > 1:
            b += g[0]-j[0]
            g[2] = 1
            j[2] = 0
            c[g[1]] += 1
            c[j[1]] -= 1
            e += 1
            f += b+e**2,
            break
print(max(f))

Submission status: Time Limit Exceeded

Solution code:
import operator
import collections

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

sushies.sort(key=operator.itemgetter(1), reverse=True)

dishes = sushies[:k]
c = collections.Counter(t for t, d in dishes)
a = sum(d for t, d in dishes)
answer = a + len(c) ** 2

for t, d in sushies[k:]:
    if t in c:
        continue

    x = -1
    for i in reversed(range(k)):
        if c[dishes[i][0]] > 1:
            x = i
            break

    if x == -1:
        break

    a = a + d - dishes[x][1]
    c[t] += 1
    c[dishes[x][0]] -= 1
    del dishes[x]
    dishes.append((t, d))

    answer = max(answer, a + len(c) ** 2)

print(answer)

Submission status: Time Limit Exceeded

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

new_arr = []
for xy in arr:
    if not xy:
        continue
    new_arr.append(xy)

new_arr.sort(reverse=True)

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

print(res)

Submission status: Time Limit Exceeded

Solution code:
import heapq

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

total = 0
types = []
remain = []
for i in range(K):
    d, t = heapq.heappop(sushi)
    total -= d
    if t not in types:
        types.append(t)
    else:
        # 余りの寿司はremainにappend
        heapq.heappush(remain, -d)

types_len = len(types)
ans = total + types_len ** 2

while sushi and remain:
    d, t = heapq.heappop(sushi)
    if t not in types:
        types.append(t)
        remain_d = heapq.heappop(remain)
        total = total - d - remain_d
        types_len += 1
        ans = max(ans, total + types_len ** 2)

print(ans)

Submission status: Time Limit Exceeded

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
        t=t[:-1]
        ans.append(g)
        retu[li[i][0]-1]+=1
print(max(ans))

Submission status: Time Limit Exceeded

Solution code:
from operator import itemgetter
n,k = map(int,input().split())
li=[]
s=[]
t=[]
g=0
ans=[]
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 li[i][0] not in s:
        s.append(li[i][0])
    else:
        t.append(li[i][1])
p=len(s)
g+=p**2
ans.append(g)
for i in range(k,n):
    if len(t)==0:
        break
    if li[i][0] not in s:
        g=g-t[-1]+li[i][1]+2*p+1
        p+=1
        ans.append(g)
        s.append(li[i][0])
print(max(ans))

Submission status: Time Limit Exceeded

Solution code:
from operator import itemgetter
n,k = map(int,input().split())
li=[]
s=[]
t=[]
g=0
ans=[]
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 li[i][0] not in s:
        s.append(li[i][0])
    else:
        t.append(li[i][1])
g+=len(s)**2
ans.append(g)
for i in range(k,n):
    if len(t)==0:
        break
    if li[i][0] not in s:
        g=g-t[-1]+li[i][1]+2*len(s)+1
        ans.append(g)
        s.append(li[i][0])
print(max(ans))

Submission status: Time Limit Exceeded

Solution code:
from itertools import product
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]
ans = 0
for pro in product(range(N), repeat=K):
    if len(set(pro)) == len(pro):
        tmp = 0
        s = []
        for p in pro: 
            s.append(td[p][0])
        tmp += len(set(s))**2
        for p in pro:
            tmp += td[p][1]
        ans = max(ans,tmp)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import queue
from collections import Counter
from operator import itemgetter

N, K = list(map(int, input().split()))
td = sorted([list(map(int, input().split())) for i in range(N)], key=itemgetter(1), reverse=True)
group = queue.PriorityQueue()
count = Counter()
ans = 0

for t, d in td[:K]:
    group.put([d, t])
    count[t] += 1
    ans += d
ans += len(count) ** 2

s = queue.PriorityQueue()
s_count = Counter()
for t, d in td[K:]:
    if not count[t] and not s_count[t]:
        s.put([-d, t])
        s_count[t] += 1

n = ans
x = len(count)
while group.qsize():
    d, t = group.get()
    if count[t] != 1:
        count[t] -= 1
        n -= (d + x**2)

        d_, t_ = s.get()
        count[t_] += 1
        x += 1
        n += (-d_ + x**2)
        ans = max(ans, n)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = list(map(int, input().split()))
tk = [list(map(int, input().split())) for i in range(n)]
ans = 0
tk.sort(key=lambda x: -x[1])
max_ = [[i, 0, 0] for i in range(n+1)]
for i in range(n):
  if max_[tk[i][0]][1]<tk[i][1]:
    max_[tk[i][0]][1], max_[tk[i][0]][2] = tk[i][1], i
max_.sort(key=lambda x: -x[1])
for i in range(1, k+1):
  neta = set()
  have = set()
  cnt = 0
  cur = 0
  for j in range(i):
    if max_[j][1]>0:
      cur += max_[j][1]
      cnt += 1
      neta.add(max_[j][0])
      have.add(max_[j][2])
  for j in range(n):
    if j in have:
      continue
    if cnt==k:
      break
    if len(neta)<i:
      neta.add(tk[j][0])
      cnt += 1
      cur += tk[j][1]
    elif len(neta)==i and tk[j][0] in neta:
      cnt += 1
      cur += tk[j][1]
  ans = max(ans, cur+len(neta)**2)
print(ans)

Submission status: Time Limit Exceeded

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)

Submission status: Time Limit Exceeded

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
for i in Rest:
    Sum=Sum+i
    Rest_s.append(Sum)
ans=0
l_r=len(Rest)
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)

Submission status: Time Limit Exceeded

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=[]
for i in sushi:
    if not i==[]:
        Max.append(max(i))
        i.remove(max(i))
        Rest=Rest+i
Max.sort(reverse=True)
Rest.sort(reverse=True)
v=len(Max)
Max_s=[]
Rest_s=[0]
Sum=0
for i in Max:
    Sum=Sum+i
    Max_s.append(Sum)
Sum=0
for i in Rest:
    Sum=Sum+i
    Rest_s.append(Sum)
ans=0
l_r=len(Rest)
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)

Submission status: Time Limit Exceeded

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:
    for sushi in sushi_array[add_index:]:
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break
        add_index += 1
    else:
        break
    for sushi in sushi_array[remove_index::-1]:
        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: Time Limit Exceeded

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:
    for sushi in sushi_array[add_index:]:
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break
        add_index += 1
    else:
        break
    for sushi in sushi_array[remove_index::-1]:
        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: Time Limit Exceeded

Solution code:
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:
    for sushi in sushi_array[add_index:]:
        if sushi_count[sushi[0]] == 0:
            kiso += sushi[1]
            sushi_count[sushi[0]] = 1
            break
        add_index += 1
    else:
        break
    for sushi in sushi_array[remove_index::-1]:
        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: Time Limit Exceeded

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

import math
import itertools

lst = list(itertools.combinations(range(N),K))
result = 0
for i in lst:
    kind = []
    oisisa = 0
    for j in range(K):
        oisisa += D[i[j]]
        if (T[i[j]] in (kind)) == False:
            kind.append(T[i[j]])
    if result < (len(kind)**2 + oisisa):
        result = len(kind)**2 + oisisa
print(result)

Submission status: Time Limit Exceeded

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)

can_use_num = sorted(can_use_num,key=lambda x: -x)

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: Time Limit Exceeded

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])
t = defaultdict(int)
x = list(range(n - 1, n - k - 1, -1))
z = [0]
for i in range(n - k, n):
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]

y = []
yd = list(t.keys())
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])
y = y[::-1]
k = 1

for k in range(1, len(y) + 1):
    for i in range(len(x) - 1, -1, -1):
        xi = x[i]
        if t[sushi[xi][0]] > 1:
            t[sushi[xi][0]] -= 1
            z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])
            x.pop(i)
            break

p = len(t.keys())
ans = 0
for i in range(len(z)):
    ans = max(ans, z[i] + (p + i) ** 2)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import copy
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])
t = {}
x = []
z = [0]
for i in range(n - 1, n - k - 1, -1):
    if sushi[i][0] in t.keys():
        x.append(i)
        t[sushi[i][0]] += 1
    else:
        t[sushi[i][0]] = 1
    z[0] += sushi[i][1]
y = []
yd = list(t.keys())
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])

for k in range(1, min(len(x), len(y)) + 1):
    xi = x[-k]
    z.append(z[-1] - sushi[xi][1] + sushi[y[k - 1]][1])

ans = 0
p = len(t.keys())
for i in range(len(z)):
    ans = max(ans, z[i] + (p + i) ** 2)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import copy
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])
t = defaultdict(int)
x = list(range(n - 1, n - k - 1, -1))
z = [0]
for i in range(n - k, n):
    t[sushi[i][0]] += 1
    z[0] += sushi[i][1]

y = []
yd = []
for i in range(n - k - 1, -1, -1):
    if sushi[i][0] not in t.keys() and sushi[i][0] not in yd:
        y.append(i)
        yd.append(sushi[i][0])
y = y[::-1]
k = 1

for k in range(1, len(y) + 1):
    for i in range(len(x) - 1, -1, -1):
        xi = x[i]
        if t[sushi[xi][0]] > 1:
            t[sushi[xi][0]] -= 1
            z.append(z[-1] - sushi[xi][1] + sushi[y[-k]][1])
            x.pop(i)
            break

p = len(t.keys())
for i in range(len(z)):
    z[i] += (p + i) ** 2
print(max(z))

Submission status: Time Limit Exceeded

Solution code:
import itertools
N,K=map(int,input().split())
ls=[]
for i in range(N):
    ls.append(list(map(int,input().split())))
ls1=list(itertools.combinations(ls,K))
ansls=[]
for ln in ls1:
    ls2=[]
    ls3=[]
    for ln1 in ln:
        ls2.append(ln1[0])
        ls3.append(ln1[1])
    set1=set(ls2)
    tsum=len(set1)**2
    dsum=sum(ls3)
    ssum=dsum+tsum
    ansls.append(ssum)
ans=max(ansls)
print(ans)

Submission status: Time Limit Exceeded

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


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

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

    mli = ml[0:K]
    mlo = deque(ml[K:])
    mli.sort(key=lambda x: x[1])

    tdic = defaultdict(int)
    for i in mli:
        tdic[i[0]] += 1

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

    print(max_score)
    return


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

if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  if len(base_t)==total_t or len(base_t)==k:
    break
  
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    elm = remains.pop(tg_idx)
    base.pop(rm_idx)
    base.append(elm)
    remains = remains[tg_idx:]
    base_t.add(elm[0])
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  if len(base_t)==total_t:
    break
  
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    elm = remains.pop(tg_idx)
    base.pop(rm_idx)
    base.append(elm)
    remains = remains[tg_idx:]
    base_t.add(elm[0])
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  if len(base_t)==k:
    break
  
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    elm = remains.pop(tg_idx)
    base.pop(rm_idx)
    base.append(elm)
    remains = remains[tg_idx:]
    base_t.add(elm[0])
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    base.pop(rm_idx)
    base.append(remains.pop(tg_idx))
    remains = remains[tg_idx:]
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
ans = 0
for i in range(total_t, 0, -1):
  t_set = set()
  k_choice = []
  duplicate_cnt = k - i
  
  for t, d in sushis:
    if t not in t_set:
      k_choice.append(d)
      t_set.add(t)
    else:
      if duplicate_cnt > 0:
        k_choice.append(d)
        duplicate_cnt -= 1
    if len(k_choice)==k:
      break
  if len(k_choice)==k:
  	ans = max(ans, sum(k_choice)+len(t_set)**2)
  else:
    continue

print(ans)

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)

ans = 0
for i in range(k, 0, -1):
  t_set = set()
  k_choice = []
  duplicate_cnt = k - i
  
  for t, d in sushis:
    if t not in t_set:
      k_choice.append(d)
      t_set.add(t)
    else:
      if duplicate_cnt > 0:
        k_choice.append(d)
        duplicate_cnt -= 1
    if len(k_choice)==k:
      break
  ans = max(ans, sum(k_choice)+len(t_set)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from itertools import combinations

n, k = map(int,input().split(' '))
sushis = [tuple(map(int, input().split(' '))) for _ in range(n)]
ans = 0

for k_choice in combinations(sushis, k):
  ans = max(ans, sum(d for t, d in k_choice) + len(set(t for t, d in k_choice))**2)
    
print(ans)

Submission status: Time Limit Exceeded

Solution code:
n,k=map(int,input().split())
a=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x: x[1])
type=set([a[i][0] for i in range(k)])
dsum=sum([x[1] for x in a[:k]])
tnum=dict()
for i in [x[0] for x in a[:k]]:
    if i not in tnum.keys():
        tnum[i]=1
    else:tnum[i]+=1
ans=dsum+len(type)**2
sushi=a[-(n-k+1)::-1]
for i in a[k:]:
    if i[0] not in type:
        type.add(i[0])
        for x in sushi[:]:
            if tnum[x[0]]>1:
                tnum[x[0]]-=1
                dsum=dsum+i[1]-x[1]
                ans=max(ans,dsum+len(type)**2)
                break
            sushi.remove(x)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
n,k=map(int,input().split())
a=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x: x[1])
type=set([a[i][0] for i in range(k)])
dsum=sum([x[1] for x in a[:k]])
tnum={i:[x[0] for x in a[:k]].count(i) for i in type}
ans=dsum+len(type)**2
sushi=a[-(n-k+1)::-1]
for i in a[k:]:
    if i[0] not in type:
        type.add(i[0])
        for x in sushi:
            if tnum[x[0]]>1:
                tnum[x[0]]-=1
                sushi.remove(x)
                dsum=dsum+i[1]-x[1]
                ans=max(ans,dsum+len(type)**2)
                break
print(ans)

Submission status: Time Limit Exceeded

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

# 美味しさ順
base_var = set([_td[0] for _td in td[:K]])
score_d = sum([_td[1] for _td in td[:K]]) + len(base_var)**2

score_list = [score_d]

candidate_list = []
for _td in td[K:]:
    if _td[0] not in base_var:
        candidate_list.append(_td)
        base_var.add(_td[0])

current = td[:K]

for candidate in candidate_list:
    current_var = set([_td[0] for _td in current])
    for i in range(len(current))[::-1]:
        _current = current[::]
        _current.pop(i)
        if len(current_var) != len(set([_td[0] for _td in _current])):
            continue
        _current.insert(0, candidate)
        var = set([_td[0] for _td in _current])
        score = sum([_td[1] for _td in _current]) + len(var) ** 2
        score_list.append(score)
        current = _current
        break

print(max(score_list))

Submission status: Time Limit Exceeded

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

n, k = map(int, input().split())
l = [list(map(int, input().split())) for _ in range(n)]

mm = 0
# very slow
for v in itertools.combinations(l, k):
    t = set()
    m = 0
    for a in v:
        t.add(a[0])
        m += a[1]
        
    m += len(t) * len(t)
    mm = max(mm, m)
print(mm)

Submission status: Time Limit Exceeded

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

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

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

Submission status: Time Limit Exceeded

Solution code:
import heapq

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

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

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

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

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

print(ans)