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:
import sys
input = sys.stdin.buffer.readline

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

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

sushi.sort()
sushi_index = n-1

fixed = {}
ate = []

for i in range(k):
    d, t = sushi[sushi_index]
    sushi_index -= 1

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

ate.sort(reverse=True)


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


score = get_score()

while sushi_index >= 0 and ate:
    d, t = sushi[sushi_index]
    sushi_index -= 1

    if t not in fixed:
        fixed[t] = d
        del ate[-1]
        score = max(score, get_score())
        if len(fixed) == k:
            break

print(score)

Submission status: Time Limit Exceeded

Solution code:
from collections import deque
n,k=map(int,input().split())
td=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x:x[1])

ans=0
kinds=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kinds:
        kinds[td[i][0]].append(td[i][1])
    else:
        kinds[td[i][0]]=deque()
        kinds[td[i][0]].append(td[i][1])
l=len(kinds)
ans+=(l**2)

ans_=ans
for i in range(k,n):
    if td[i][0] not in kinds:
        mi=[1000000000,1000000000]
        for j in kinds:
            if len(kinds[j])>1:
                if mi[1]>kinds[j][-1]:
                    mi=[j,kinds[j][-1]]
        if mi==[1000000000,1000000000]:
            break
        else:
            ans=ans-l**2+(l+1)**2-mi[1]+td[i][1]
            kinds[mi[0]].pop()
            kinds[td[i][0]]=deque()
            kinds[td[i][0]].append(td[i][1])
            ans_=max(ans,ans_)
            l+=1

print(ans_)

Submission status: Time Limit Exceeded

Solution code:
from collections import deque
n,k=map(int,input().split())
td=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x:x[1])

ans=0
kinds=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kinds:
        kinds[td[i][0]].append(td[i][1])
    else:
        kinds[td[i][0]]=deque()
        kinds[td[i][0]].append(td[i][1])
ans+=(len(kinds)**2)

ans_=[ans]
for i in range(k,n):
    if td[i][0] not in kinds:
        mi=[1000000000,1000000000]
        for j in kinds:
            if len(kinds[j])>1:
                if mi[1]>kinds[j][-1]:
                    mi=[j,kinds[j][-1]]
        if mi==[1000000000,1000000000]:
            break
        else:
            ans=ans-len(kinds)**2+(len(kinds)+1)**2-mi[1]+td[i][1]
            kinds[mi[0]].pop()
            kinds[td[i][0]]=deque()
            kinds[td[i][0]].append(td[i][1])
            ans_.append(ans)

print(max(ans_))

Submission status: Time Limit Exceeded

Solution code:
n,k=map(int,input().split())
td=sorted([list(map(int,input().split())) for i in range(n)],reverse=True,key=lambda x:x[1])
#print(td)
ans=0
kinds=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kinds:
        kinds[td[i][0]].append(td[i][1])
    else:
        kinds[td[i][0]]=[td[i][1]]
ans+=(len(kinds)**2)
#kindsの中はソートされている状況
#print(kinds)
ans_=[ans]
for i in range(k,n):
    if td[i][0] not in kinds:
        mi=[1000000000,1000000000]
        for j in kinds:
            if len(kinds[j])>1:
                if mi[1]>kinds[j][-1]:
                    mi=[j,kinds[j][-1]]
        if mi==[1000000000,1000000000]:
            break
        else:
            ans=ans-len(kinds)**2+(len(kinds)+1)**2-mi[1]+td[i][1]
            kinds[mi[0]].pop(-1)
            kinds[td[i][0]]=[td[i][1]]
            ans_.append(ans)
#print(kinds)
print(max(ans_))

Submission status: Time Limit Exceeded

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()

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

sushi = [[] for _ in range(N)]

for i in range(N):
    a,b = map(int, input().split())
    a -= 1
    sushi[a].append(b)

first = []
second = []
for a in sushi:
    if a:
        sort_a = sorted(a, reverse=True)
        flg = 1
        for num in sort_a:
            if flg:
                first.append((num, 1))
                flg = 0
            else:
                second.append((num, 0))
first.sort(reverse=True)                
second.sort(reverse=True)                

i = 0
j = 0
ans = 0

for i in range(len(first)):
    sum_first = 0
    sum_second = 0
    cnt = 0
    if i > K-1:
        break
    for j in range(i+1):
        sum_first += first[j][0]
        cnt += 1
    for k in range(K-i-1):
        if k > len(second) - 1:
            break
        sum_second += second[k][0]
  
    ans = max(ans, sum_first + sum_second + cnt*cnt)

print(ans)

Submission status: Time Limit Exceeded

Solution code:

from collections import deque
n,k= map(int,input().split())
td=[]
for i in range(n):
    td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
select=[]
noselect=[]
cat=[]
for tdi in td:
    if tdi[0] in cat:
        noselect.append(tdi)
        continue
    select.append(tdi)
    cat.append(tdi[0])
    if len(select)==k:
        break
i=0
while len(select) < k:
    select.append(noselect[i])
    i+=1
noselect=noselect[i:]

select.sort(key=lambda x:x[1])
select=deque(select)
#cat=list(set([tdi[0] for tdi in select]))
for no in noselect:
    if no[0] not in cat:
        continue
    if (2 * len(cat)) - 1 < no[1] - select[0][1]:
        select.popleft()
        select.append(no)
        cat=list(set([tdi[0] for tdi in select]))
    else:
        break
print(sum([tdi[1] for tdi in select])+(len(cat)**2))

Submission status: Time Limit Exceeded

Solution code:

from collections import deque
n,k= map(int,input().split())
td=[]
for i in range(n):
    td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
select=[]
noselect=[]
cat=[]
for tdi in td:
    if tdi[0] in cat:
        noselect.append(tdi)
        continue
    select.append(tdi)
    cat.append(tdi[0])
    if len(select)==k:
        break
i=0
noselect=deque(noselect)
while len(select) < k:
    select.append(noselect.popleft())
    i+=1
select.sort(key=lambda x:x[1])
select=deque(select)
cat=list(set([tdi[0] for tdi in select]))
for no in noselect:
    if no[0] not in cat:
        continue
    if (2 * len(cat)) - 1 < no[1] - select[0][1]:
        select.popleft()
        select.append(no)
        cat=list(set([tdi[0] for tdi in select]))
    else:
        break
print(sum([tdi[1] for tdi in select])+(len(cat)**2))

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict

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

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

neta = set()
first = []
second = []
for s in sushi:
    t, d = s
    if t not in neta:
        first.append(d)
        neta.add(t)
    else:
        second.append(d)

first.sort(reverse=True)
second.sort(reverse=True)

kinds = min(k, len(first))
ans = 0

for i in range(1, kinds + 1):
    point = sum(first[:i])
    point += sum(second[:k - i])
    point += i ** 2
    ans = max(ans, point)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    other_sushi += value

most_tasty.sort(reverse=True)
other_sushi.sort(reverse=True)
for i in range(1, k + 1):
    if i > len(most_tasty):
        break
    num = min(i, len(most_tasty))
    point = sum(most_tasty[:i])
    point += sum(other_sushi[:(k - i)])
    point += i ** 2
    ans = max(ans, point)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    other_sushi += value

most_tasty.sort(reverse=True)
other_sushi.sort(reverse=True)
for i in range(1, k + 1):
    num = min(i, len(most_tasty))
    point = sum(most_tasty[:num])
    point += sum(other_sushi[:(k - num)])
    point += num ** 2
    ans = max(ans, point)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    for v in value:
        other_sushi.append(v)

most_tasty.sort(reverse=True)
other_sushi.sort(reverse=True)
for i in range(1, k + 1):
    num = min(i, len(most_tasty))
    point = sum(most_tasty[:num])
    point += sum(other_sushi[:(k - num)])
    point += num ** 2
    ans = max(ans, point)

print(ans)

Submission status: Time Limit Exceeded

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

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

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

used = [0 for _ in range(N)]
dup = []
deli = 0

for i in range(K):
    t, d = sushi[i]
    if used[t - 1]:
        dup.append(d)
    used[t - 1] = 1
    deli += d

sat = sum(used)**2 + deli

dup.sort(reverse=True)

for i in range(N - K):
    if not dup:
        break
    t, d = sushi[K + i]
    if used[t - 1]:
        continue
    used[t - 1] = 1
    deli -= dup.pop()
    deli += d
    sat = max(sat, sum(used)**2 + deli)

print(sat)

Submission status: Time Limit Exceeded

Solution code:
import sys
import heapq as h

input=sys.stdin.readline

N,K=map(int,input().split())
sushi=[[0,0] for i in range(0,N)]
for i in range(0,N):
    t,k=map(int,input().split())
    sushi[i][1]=t
    sushi[i][0]=k

sushikind=[[] for i in range(0,N+1)]
for i in range(0,N):
    sushikind[sushi[i][1]].append(sushi[i][0])

kindmax=[0 for i in range(0,N+1)]

for i in range(0,N):
    if sushikind[i]:
        kindmax[i]=max(sushikind[i])

sushi.sort(reverse=True)
sum=0
kind=[]
kind=set(kind)
extra=[]
need=[]
for i in range(0,K):
    sum+=sushi[i][0]
    kind.add(sushi[i][1])
    if sushi[i][0]==kindmax[sushi[i][1]]:
        if sushi[i] in need:
            extra.append(sushi[i][0])
        else:
            need.append(sushi[i])
    else:
        extra.append(sushi[i][0])


set={i for i in range(1,N+1)}
set-=kind
appendix=[kindmax[i] for i in set if kindmax[i]!=0 ]
extra.sort()
appendix.sort(reverse=True)
k=len(kind)
first=k
ans=k**2+sum
for i in range(k,min(K,k+len(appendix))):
    k+=1
    sum+=appendix[k-first-1]-extra[k-first-1]
    test=k**2+sum
    if test>ans:
        ans=test

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import sys
import heapq as h

input=sys.stdin.readline

N,K=map(int,input().split())
sushi=[[0,0] for i in range(0,N)]
for i in range(0,N):
    t,k=map(int,input().split())
    sushi[i][1]=t
    sushi[i][0]=k

sushikind=[[] for i in range(0,N+1)]
for i in range(0,N):
    sushikind[sushi[i][1]].append(sushi[i][0])

kindmax=[0 for i in range(0,N+1)]

for i in range(0,N):
    if sushikind[i]:
        kindmax[i]=max(sushikind[i])

sushi.sort(reverse=True)
sum=0
kind=[]
kind=set(kind)
extra=[]
need=[]
for i in range(0,K):
    sum+=sushi[i][0]
    kind.add(sushi[i][1])
    if sushi[i][0]==kindmax[sushi[i][1]]:
        if sushi[i] in need:
            extra.append(sushi[i][0])
        else:
            need.append(sushi[i])
    else:
        extra.append(sushi[i][0])


appendix=[kindmax[i] for i in range(0,N) if kindmax[i]!=0 and (not i in kind)]
extra.sort()
appendix.sort(reverse=True)
k=len(kind)
first=k
ans=k**2+sum
for i in range(k,min(K,k+len(appendix))):
    k+=1
    sum+=appendix[k-first-1]-extra[k-first-1]
    test=k**2+sum
    if test>ans:
        ans=test

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import sys
import heapq as h

input=sys.stdin.readline

N,K=map(int,input().split())
sushi=[[0,0] for i in range(0,N)]
for i in range(0,N):
    t,k=map(int,input().split())
    sushi[i][1]=t
    sushi[i][0]=k

sushikind=[[] for i in range(0,N+1)]
for i in range(0,N):
    sushikind[sushi[i][1]].append(sushi[i][0])

kindmax=[0 for i in range(0,N+1)]

for i in range(0,N):
    if sushikind[i]:
        kindmax[i]=max(sushikind[i])

sushi.sort(reverse=True)
sum=0
kind=[]
kind=set(kind)
extra=[]
h.heapify(extra)
need=[]
for i in range(0,K):
    sum+=sushi[i][0]
    kind.add(sushi[i][1])
    if sushi[i][0]==kindmax[sushi[i][1]]:
        if sushi[i] in need:
            h.heappush(extra,sushi[i][0])
        else:
            need.append(sushi[i])
    else:
        h.heappush(extra,sushi[i][0])


appendix=[-kindmax[i] for i in range(0,N) if kindmax[i]!=0 and (not i in kind)]
h.heapify(appendix)

k=len(kind)
ans=k**2+sum
for i in range(k,min(K,k+len(appendix))):
    k+=1
    sum-=h.heappop(appendix)+h.heappop(extra)
    test=k**2+sum
    if test>ans:
        ans=test

print(ans)

Submission status: Time Limit Exceeded

Solution code:
from itertools import combinations

N, K = list(map(int,input().split()))
t, d = [0]*(N+1), [0]*(N+1)
for i in range(N):
  t[i+1], d[i+1] = map(int,input().split())

l = [i for i in range(1, N+1)]
ans = 0
for v in combinations(l, K):
  bp, kp = 0, 0
  s = set()
  for k in range(K):
    bp += d[v[k]]
    s.add(t[v[k]])
  kp = len(s)*len(s)
  ans = max(ans,bp+kp)
print(ans)

Submission status: Time Limit Exceeded

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

list0 = []
list1 = []
set1 = set()

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

td.sort(key=lambda x: (x[0], -x[1]))

for t, d in td:
    if t in set1:
        list0.append(d)
    else:
        set1.add(t)
        list1.append(d)

list0.sort(reverse=True)
list1.sort(reverse=True)

various = min(K, len(list1))
selected0 = []
selected1 = []
for i in range(various):
    selected1.append(list1.pop(0))
if K > various:
    for i in range(K - various):
        selected0.append(list0.pop(0))


ans = sum(selected0) + sum(selected1) + various ** 2

mn = selected1[-1]
while list0 and list0[0] > mn:
    mn = selected1.pop(-1)
    selected0.append(list0.pop(0))
    various -= 1
    tmp = sum(selected0) + sum(selected1) + various ** 2
    if tmp <= ans:
        break
    ans = tmp

print(ans)

Submission status: Time Limit Exceeded

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

list0 = []
list1 = []
set1 = set()

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

td.sort(key=lambda x: (x[0], -x[1]))

for t, d in td:
    if t in set1:
        list0.append(d)
    else:
        set1.add(t)
        list1.append(d)

list0.sort(reverse=True)
list1.sort(reverse=True)

various = min(K, len(list1))
selected0 = []
selected1 = []
for i in range(various):
    selected1.append(list1.pop(0))
if K > various:
    for i in range(K - various):
        selected0.append(list0.pop(0))


ans = sum(selected0) + sum(selected1) + various ** 2

while list0 and list0[0] > selected1[-1]:
    selected1.pop(-1)
    selected0.append(list0.pop(0))
    various -= 1
    tmp = sum(selected0) + sum(selected1) + various ** 2
    if tmp <= ans:
        break
    ans = tmp

print(ans)

Submission status: Time Limit Exceeded

Solution code:

import heapq

n,k,*td=map(int,open(0).read().split())
td = [[t,d] for t,d in zip(td[::2],td[1::2])]
td = sorted(td,key = lambda x:x[1],reverse=True)

score = [0] * k
taste_sum = 0
extra_taste = []
eat_num = 0
eat_kind = set()

for t,d in td:
    if t in eat_kind:
        if eat_num == k:
            continue
        else:
            heapq.heappush(extra_taste, d)
            eat_num += 1
    else:
        eat_kind.add(t)
        taste_sum += d
        eat_num += 1
        if eat_num > k:
            eat_num -= 1
            heapq.heappop(extra_taste)
    
    if eat_num == k:
        score[len(eat_kind)-1] = len(eat_kind)**2 + taste_sum + sum(extra_taste)
            
    if len(eat_kind) == k:
        break
print(max(score))

Submission status: Time Limit Exceeded

Solution code:
from bisect import insort
n,k,*td=map(int,open(0).read().split())
td = [[t,d] for t,d in zip(td[::2],td[1::2])]
td = sorted(td,key = lambda x:x[1],reverse=True)

score = [0] * k
taste_sum = 0
extra_taste = []
eat_num = 0
eat_kind = set()

for t,d in td:
    if t in eat_kind:
        if eat_num == k:
            continue
        else:
            insort(extra_taste, d)
            eat_num += 1
    else:
        eat_kind.add(t)
        taste_sum += d
        eat_num += 1
        if eat_num > k:
            eat_num -= 1
            extra_taste = extra_taste[1:]
    
    if eat_num == k:
        score[len(eat_kind)-1] = len(eat_kind)**2 + taste_sum + sum(extra_taste)
            
    if len(eat_kind) == k:
        break
print(max(score))

Submission status: Time Limit Exceeded

Solution code:
from bisect import insort
n,k,*td=map(int,open(0).read().split())
td = [[t,d] for t,d in zip(td[::2],td[1::2])]
td = sorted(td,key = lambda x:x[1],reverse=True)

score = [0] * k
taste_sum = 0
extra_taste = []
eat_num = 0
eat_kind = set()

for t,d in td:
    if t in eat_kind:
        insort(extra_taste, d)
        if eat_num == k:
            extra_taste = extra_taste[1:]
            continue
        else:
            eat_num += 1
    else:
        eat_kind.add(t)
        taste_sum += d
        eat_num += 1
        if eat_num > k:
            eat_num -= 1
            extra_taste = extra_taste[1:]
    
    if eat_num == k:
        score[len(eat_kind)-1] = len(eat_kind)**2 + taste_sum + sum(extra_taste)
            
    if len(eat_kind) == k:
        break
print(max(score))

Submission status: Time Limit Exceeded

Solution code:
N,K = map(int,input().split())
ans = [0 for i in range(N+1)]
food = {}
food1 = []
for i in range(1,N+1):
    food[i] = 0
for i in range(N):
    t,d = map(int,input().split())
    food1.append((d,t))
food1.sort()
syurui = 0
sum = 0
x = []
for i in range(1,N+1):
    d,t = food1[-i]
    if food[t] == 0:
        if i > K:
            if len(x) == 0:
                break
        sum += d
        food[t] = 1
        syurui += 1
        if i > K:
            y = x.pop(0)
            sum -= y
    else:
        if i <= K:
            sum += d
            x.append(d)
            x.sort()
    if syurui == K:
        ans[K] = max(sum + K**2,ans[K])
        break
    else:
        ans[syurui] = max(sum + syurui**2,ans[syurui])
print(max(ans))

Submission status: Time Limit Exceeded

Solution code:
def p_d():
    from collections import defaultdict
    from heapq import heappush, heappop
    N, K = map(int, input().split())
    li = []
    for _ in range(N):
        t, d = map(int, input().split())
        li.append((t, d))
    li.sort(key=lambda x: x[1], reverse=True)
    sushi = defaultdict(list)
    ans = 0
    for i in range(K):
        heappush(sushi[li[i][0]], li[i][1])
        ans += li[i][1]
    ans += len(sushi.keys()) ** 2

    ans_list = [ans]
    while True:
        k, v = -1, 10 ** 9 + 1
        for key in sushi.keys():
            if len(sushi[key]) > 1:
                tmp = heappop(sushi[key])
                if v > tmp:
                    if k != -1:
                        heappush(sushi[k], v)
                    k, v = key, tmp
                else:
                    heappush(sushi[key], tmp)
        if k == -1:
            break
        while K < N and li[K][0] in sushi.keys():
            K += 1
        if K == N:
            break
        ans -= v
        ans -= len(sushi.keys()) ** 2
        heappush(sushi[li[K][0]], li[K][1])
        ans += li[K][1]
        ans += len(sushi.keys()) ** 2
        ans_list.append(ans)
    print(max(ans_list))


if __name__ == '__main__':
    p_d()

Submission status: Time Limit Exceeded

Solution code:
import itertools
import collections
N,K = map(int,input().split())
Sushi = [list(map(int,input().split(" "))) for i in range(N)]
Com = list(itertools.combinations(list(range(N)),K))
Max = 0
for index in Com:
    v = []
    point = 0
    for n in list(index):
        point += Sushi[n][1]
        v.append(Sushi[n][0])
    vary = len(list(set(v)))
    point += vary**2
    Max = max(Max,point)
print(Max)

Submission status: Time Limit Exceeded

Solution code:
# -*- coding: utf-8 -*-

import sys
from heapq import heapify, heappush, heappop
from collections import Counter

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10 ** 9)
INF = 10 ** 18
MOD = 10 ** 9 + 7

N, K = MAP()
DT = []
tmp = [[] for i in range(N)]
for i in range(N):
    t, d = MAP()
    DT.append((d, t))
DT.sort(reverse=1)

que = []
C = Counter()
dcnt = 0
# ポイント上位K個をまず取る
for i in range(K):
    d, t = DT[i]
    dcnt += d
    C[t] += 1
    heappush(que, (d, t))
heapify(que)
tcnt = len(C)

ans = dcnt + tcnt**2
for i in range(K, N):
    d2, t2 = DT[i]
    # 既に持っている種類ならスキップ
    if C[t2] > 0:
        continue
    tmp = []
    while que:
        d, t = heappop(que)
        # 2つ以上ある種類が来たら交換候補
        if C[t] >= 2:
            break
        else:
            tmp.append((d, t))
    # 全部1種類ずつになったらもう種類数を増やせないので終了
    if not que:
        break
    # この候補と交換してみる
    dcnt -= d
    dcnt += d2
    C[t] -= 1
    C[t2] += 1
    tcnt += 1
    ans = max(ans, dcnt + tcnt**2)
    # 今回使った分を含め、キューを元に戻す
    heappush(que, (d2, t2))
    for d, t in tmp:
        heappush(que, (d, t))
print(ans)

Submission status: Time Limit Exceeded

Solution code:
def get_point(lst):
	types_point = len(set([i[0] for i in lst])) ** 2
	satis_point = sum([i[1] for i in lst])
	return types_point + satis_point

n, k = [int(a) for a in input().split()]
lst = []
for i in range(n):
	t, s = [int(a) for a in input().split()]
	lst.append([t, s])

for i in range(n - k):
	tmp_list = []
	for j in range(len(lst)):
		tmp_satis = lst[0 : j] + lst[j + 1 : ]
		tmp_total = get_point(tmp_satis)
		tmp_list.append(tmp_total)
	tmp_max = max(tmp_list)
	delete_idx = tmp_list.index(tmp_max)  #貢献が小さいidx
	lst = lst[0 : delete_idx] + lst[delete_idx + 1 : ]
print (get_point(lst))

Submission status: Time Limit Exceeded

Solution code:
n, k = [int(a) for a in input().split()]
types = []
satis = []
for i in range(n):
	t, s = [int(a) for a in input().split()]
	types.append(t)
	satis.append(s)
total = len(set(types)) ** 2 + sum(satis)

for i in range(n - k):
	tmp_list = []
	for j in range(len(types)):
		tmp_types = types[0 : j] + types[j + 1 : ]
		tmp_satis = satis[0 : j] + satis[j + 1 : ]
		tmp_total = len(set(tmp_types)) ** 2 + sum(tmp_satis)
		tmp_list.append(tmp_total)
	tmp_max = max(tmp_list)
	delete_idx = tmp_list.index(tmp_max)  #貢献が小さいidx
	types = types[0 : delete_idx] + types[delete_idx + 1 : ]
	satis = satis[0 : delete_idx] + satis[delete_idx + 1 : ]
print (len(set(types)) ** 2 + sum(satis))

Submission status: Time Limit Exceeded

Solution code:
import sys
sys.setrecursionlimit(1000000)
n, k = list(map(int, input().split(' ')))
sushi = {}
for i in range(n):
    t, d = list(map(int, input().split(' ')))
    sushi[t] = sushi.get(t, []) + [d]
for key in sushi:
    sushi[key] = sorted(sushi[key])[::-1]

memo = {}
used = [False] * 100001
def eat(sushi, ct, pt):
    hoge = str(sorted(sushi.items()))
    if hoge in memo:
        return memo[hoge]

    # 終わり
    if ct == k:
        # print (pt, sum([v for _, v in used.items()]) ** 2)
        return pt + sum(used) ** 2
    same = 0
    diff = 0

    # 同じやつ
    target = None
    max_val = -1
    for key in sushi:
        if used[key]:
            if sushi[key] and max_val < sushi[key][0]:
                target = key
                max_val = sushi[key][0]
    if target is not None:
        val = sushi[target].pop(0)
        same = eat(sushi, ct+1, pt+val)
        # 戻す
        sushi[target].insert(0, val)

    # 新しいやつ
    target = None
    max_val = -1
    for key in sushi:
        if not used[key]:
            if sushi[key] and max_val < sushi[key][0]:
                target = key
                max_val = sushi[key][0]
    if target is not None:
        val = sushi[target].pop(0)
        used[target] = True
        diff = eat(sushi, ct+1, pt+val)
        # 戻す
        used[target] = False
        sushi[target].insert(0, val)
    memo[hoge] = max(same, diff)
    return memo[hoge]

print (eat(sushi, 0, 0))

Submission status: Time Limit Exceeded

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

l=[] # 寿司リスト。ただし，「おいしさ」を第1要素にする
dj=[] # 寿司タプルをそのまま入れる
s=set()

for i in range(N):
    t,d=map(int,input().split())
    l.append((d,t))
l.sort(reverse=True) 

ans=K**2
j=0 # 既に食べたことのある種類の寿司を更に食べた回数

for i in range(K): # i番目に食べる寿司について
    while l and l[0][1] in s:
        dj.append(l.pop(0))
    if (l and dj and l[0][0]>=dj[0][0]-(2*K-1-2*j)) or not dj:
        ans+=l[0][0]
        s.add(l[0][1])
        l.pop(0)
    else:
        ans+=dj[0][0]-(2*K-1-2*j)
        dj.pop(0)
        j+=1

print(ans)

Submission status: Time Limit Exceeded

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

t = [0]*N
d = [0]*N
for i in range(N):
    t[i], d[i] = map(int, input().split())
ans = int(0)
s = max(t)
value = [0]*s
kind = int(0)
k = []

top = [0]*s
all = []
delcan = []
pairs = zip(t,d)

all = sorted(pairs, key=lambda x: x[1],reverse=True)
for i in range(K):
    value[0] += all[i][1]
    if not all[i][0] in k:
        k.append(all[i][0])
        kind += 1
    else:
        delcan.append(all[i][1])

for i in range(N):
    if top[t[i]-1]<d[i]:
        top[t[i]-1] = d[i]

for i in k:
    top[i-1] = 0
while 0 in top:
    top.remove(0)
top.sort(reverse=True)
delcan.sort()

value[0] += kind**2
ans = value[0]
for i in range(min(len(top),len(delcan))):
    if top and delcan:
        value[i+1] = value[i] - delcan[i] + top[i] + (kind * 2 + 1)
    kind +=1
    if ans<value[i+1]:
        ans = value[i+1]
print(ans)

Submission status: Time Limit Exceeded

Solution code:
from heapq import heappop, heappush

N, K = map(int, input().split())
td = sorted([list(map(int, input().split()))
             for _ in range(N)], key=lambda x: x[1], reverse=True)
q = []
v = []
cnt = 0
for t, d in td[:K]:
    cnt += d
    if t in v:
        heappush(q, d)
    else:
        v.append(t)
cnt += len(v) ** 2
ans = cnt
for t, d in td[K:]:
    if t not in v and len(q) != 0:
        x = heappop(q)
        cnt += d + 2 * len(v) + 1 - x
        v.append(t)
        ans = max(ans, cnt)
print(ans)

Submission status: Time Limit Exceeded

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

N, K = map(int,input().split())
td = [list(map(int, input().split())) for _ in range(N)]
td.sort(key=lambda x: (x[0],x[1]),reverse=True)

x= []
y= []
y.append(td[0][1])
for i in range(1,N):
    if td[i][0] != td[i-1][0]:
        y.append(td[i][1])
    else:
        x.append(td[i][1])
x.sort(reverse=True)
y.sort(reverse=True)
ans = 0
for i in range(1,min(K,len(y))+1):
    ans = max(ans, sum(y[:i])+sum(x[:max(0,K-i)])+i*i)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
# d
I = lambda: map(int, input().split())
n, k = I()
a = []
for _ in range(n):
    t, d = I()
    a += [d, t],
a.sort(key = lambda x: x[0], reverse = True)
b = 0
c = [0]*(n+1)
e = 0
s = []
for i in a[:k]:
    b += i[0]
    if c[i[1]] == 0: e += 1
    c[i[1]] += 1
    s += i,
f = b+e**2
for i in range(k, n):
    g = a[i]
    if c[g[1]] > 0: continue
    for j in range(k):
        h = s[::-1][j]
        if c[h[1]] > 1:
            b += g[0]-h[0]
            c[g[1]] += 1
            c[h[1]] -= 1
            s.pop(j)
            s += g,
            e += 1
            f = max(f, b+e**2)
            break
print(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],
a.sort(key = lambda a: a[0], reverse = True)
b = 0
c = [0]*(n+1)
e = 0
s = []
for i in a[:k]:
    b += i[0]
    if c[i[1]] == 0: e += 1
    c[i[1]] += 1
    s += i,
f = b+e**2
for i in range(k, n):
    g = a[i]
    if c[g[1]] > 0: continue
    for j in range(k):
        h = s[::-1][j]
        if c[h[1]] > 1:
            b += g[0]-h[0]
            c[g[1]] += 1
            c[h[1]] -= 1
            s.pop(j)
            s += g,
            e += 1
            f = max(f, b+e**2)
            break
print(f)

Submission status: Time Limit Exceeded

Solution code:
# d
I = lambda: map(int, input().split())
n, k = I()
a = []
for _ in range(n):
    t, d = I()
    a += [d, t, 0],
a.sort(key = lambda a: a[0], 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 = max(f, b+e**2)
            break
print(f)

Submission status: Time Limit Exceeded

Solution code:
# d
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))