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 heapq import heapify,heappush,heappop
N,K=map(int,input().split())
sushi=[]
for i in range(N):
  t,d=map(int,input().split())
  sushi.append((d,t-1))
sushi.sort(key=lambda x:-x[0])
inf=float("inf")
oishisa=[-inf]*(N+1)
neta=[0]*N
heaps=[[] for i in range(N)]
S=[]
heapify(S)
for i in range(N):
  heapify(heaps[i])
a,k=0,0

for d,t in sushi[:K]:
  a+=d
  neta[t]+=1
  if neta[t]==1:
    k+=1
  heappush(heaps[t],d)

oishisa[k]=a
for i in range(N):
  if neta[i]>=2:
    d=heappop(heaps[i])
    heappush(S,(d,i))
for d,t in sushi[K:]:
  if len(S)==0:
    break
  if neta[t]==0:
    d_,t_=heappop(S)
    neta[t_]-=1
    if neta[t_]>=2:
      d__=heappop(heaps[t_])
      heappush(S,(d__,t_))
    heappush(heaps[t],d)
    oishisa[k+1]=oishisa[k]+d-d_
    k+=1
ans=0
for i in range(N+1):
  ans=max(oishisa[i]+i*i,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

def inpl(): return [int(i) for i in input().split()]

N,K = inpl()
d = {}
ans = 0
n0 = []
n1 = []

for i in range(N):
    ti,di = inpl()
    if (ti in d.keys()):
        d[ti].append(di)
    else:
        d[ti] = [di]

for i in d.keys():
    sortedd = sorted(d[i],reverse = True)
    n1.append(sortedd[0])
    n0 += sortedd[1:]

n0.sort(reverse = True)
n1.sort(reverse = True)

for i in range(K+1):
    if(K-i <= len(n0) and 0<=K-i):
        total = sum(n1[:i])+sum(n0[:K-i])+(i)**2
        if(ans<total):
            ans = total

print(ans)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

def inpl(): return [int(i) for i in input().split()]

N,K = inpl()
d = {}
ans = 0
n0 = []
n1 = []

for i in range(N):
    ti,di = inpl()
    if (ti in d.keys()):
        d[ti].append(di)
    else:
        d[ti] = [di]

for i in d.keys():
    sortedd = sorted(d[i])
    sortedd.reverse()
    n1.append(d[i][0])
    n0 += d[i][1:]

n0.sort()
n0.reverse()
n1.sort()
n1.reverse()

for i in range(K+1):
    if(K-i <= len(n0) and 0<=K-i):
        total = sum(n1[:i])+sum(n0[:K-i])+(i)**2
        if(ans<total):
            ans = total

print(ans)

Submission status: Wrong Answer

Solution code:
from math import pow
from collections import deque, Counter
import sys
stdin = sys.stdin
def li(): return map(int, stdin.readline().split())

def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    # logger.debug(target_queue)
    compare_queue = deque(sorted_recs[K:])
    count = Counter()
    for neta_point in target_queue:
        count[neta_point[0]] += 1
    # logger.debug(count)

    cardinality = len(count.keys())
    sum_target_points = sum([n_p[1] for n_p in target_queue])
    score = pow(cardinality, 2) + sum_target_points
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if count[comp[0]] > 0:
            continue
        while target_queue:
            old = target_queue.pop()
            if count[old[0]] < 2:
                # swap this neta will not increase cardinality but lower neta point
                continue
            else: # check if score will go up
                inc = 2 * cardinality + 1
                dec = old[1] - comp[1]
                delta = inc - dec
                # logger.debug(f"delta is {delta}")
                # if delta >= 0:
                #     logger.debug("SOME")
                #     logger.debug(f"{old} out")
                #     logger.debug(f"{comp} in")
                #     cardinality += 1
                #     score += delta
                #     count[old[0]] -= 1
                #     count[comp[0]] +=1
                # else: #increase cardinality will bring more lost than gain
                #     # exit here
                #     return score
                new_score = new_scores[-1] + delta
                new_scores.append(new_score)

                cardinality += 1
                count[old[0]] -= 1
                count[comp[0]] +=1
    # both queues are drained
    return max(new_scores)
N, K = li()
records = [tuple(li()) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Wrong Answer

Solution code:
from math import pow
from collections import deque, Counter
# from loguru import logger


def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    # logger.debug(target_queue)
    compare_queue = deque(sorted_recs[K:])
    count = Counter()
    for neta_point in target_queue:
        count[neta_point[0]] += 1
    # logger.debug(count)

    cardinality = len(count.keys())
    sum_target_points = sum([n_p[1] for n_p in target_queue])
    score = pow(cardinality, 2) + sum_target_points
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if count[comp[0]] > 0:
            continue
        while target_queue:
            old = target_queue.pop()
            if count[old[0]] < 2:
                # swap this neta will not increase cardinality but lower neta point
                continue
            else: # check if score will go up
                inc = 2 * cardinality + 1
                dec = old[1] - comp[1]
                delta = inc - dec
                # logger.debug(f"delta is {delta}")
                # if delta >= 0:
                #     logger.debug("SOME")
                #     logger.debug(f"{old} out")
                #     logger.debug(f"{comp} in")
                #     cardinality += 1
                #     score += delta
                #     count[old[0]] -= 1
                #     count[comp[0]] +=1
                # else: #increase cardinality will bring more lost than gain
                #     # exit here
                #     return score
                new_score = new_scores[-1] + delta
                new_scores.append(new_score)

                cardinality += 1
                count[old[0]] -= 1
                count[comp[0]] +=1
    # both queues are drained
    return max(new_scores)
N, K = map(int, input().split())
records = [tuple(map(int, input().split())) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Wrong Answer

Solution code:
from math import pow
from collections import deque, Counter
# from loguru import logger


def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    # logger.debug(target_queue)
    compare_queue = deque(sorted_recs[K:])
    count = Counter()
    for neta_point in target_queue:
        count[neta_point[0]] += 1
    # logger.debug(count)

    cardinality = len(count.keys())
    sum_target_points = sum([n_p[1] for n_p in target_queue])
    score = pow(cardinality, 2) + sum_target_points
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in count:
            continue
        while target_queue:
            old = target_queue.pop()
            if count[old[0]] < 2:
                # swap this neta will not increase cardinality but lower neta point
                continue
            else: # check if score will go up
                inc = 2 * cardinality + 1
                dec = old[1] - comp[1]
                delta = inc - dec
                # logger.debug(f"delta is {delta}")
                # if delta >= 0:
                #     logger.debug("SOME")
                #     logger.debug(f"{old} out")
                #     logger.debug(f"{comp} in")
                #     cardinality += 1
                #     score += delta
                #     count[old[0]] -= 1
                #     count[comp[0]] +=1
                # else: #increase cardinality will bring more lost than gain
                #     # exit here
                #     return score
                new_score = new_scores[-1] + delta
                new_scores.append(new_score)

                cardinality += 1
                count[old[0]] -= 1
                count[comp[0]] +=1
    # both queues are drained
    return max(new_scores)
N, K = map(int, input().split())
records = [tuple(map(int, input().split())) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Wrong Answer

Solution code:

def read_input():
    n, k = map(int, input().split())
    sushi = []

    for i in range(n):
        sushi.append(list(map(int, input().split())))

    return n, k, sushi

def check_value(sushi, sushi_list):
    oishii = 0
    neta = []

    for idx in sushi_list:
        oishii += sushi[idx][1]
        neta.append(sushi[idx][0])

    neta = set(neta)
    neta_var = len(neta)

    return oishii + (neta_var * neta_var)



if __name__ == '__main__':
    n, k, sushi = read_input()
    dp_ = dict()
    sushi_eat_ = dict()

    for j in range(k):
        dp_[(0, j)] = 0
        sushi_eat_[(0, j)] = []


    for i in range(n):
        for j in range(k):
            t1 = dp_[(i, j)]
            sushi_eat_temp = sushi_eat_[(i, j)].copy()
            sushi_eat_temp.append(i)
            t2 = check_value(sushi, sushi_eat_temp)

            if len(sushi_eat_[(i, j)]) < k:
                if t1 >= t2:
                    dp_[(i + 1, j)] = t1
                    sushi_eat_[(i + 1, j)] = sushi_eat_[(i, j)].copy()
                else:
                    dp_[(i + 1, j)] = t2
                    sushi_eat_[(i + 1, j)] = sushi_eat_temp
            else:
                dp_[(i + 1, j)] = t1
                sushi_eat_[(i + 1, j)] = sushi_eat_[(i, j)].copy()


    print(max(dp_.values()))

Submission status: Wrong Answer

Solution code:
import sys
stdin = sys.stdin
 
sys.setrecursionlimit(10**5) 
 
def li(): return map(int, stdin.readline().split())
def li_(): return map(lambda x: int(x)-1, stdin.readline().split())
def lf(): return map(float, stdin.readline().split())
def ls(): return stdin.readline().split()
def ns(): return stdin.readline().rstrip()
def lc(): return list(ns())
def ni(): return int(stdin.readline())
def nf(): return float(stdin.readline())

from heapq import heappush, heappop
from collections import deque

n,k = li()
td = [tuple(li()) for _ in range(n)]

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

kindnum = 0
kindset = set()
hque = []

ans = 0

for ti, di in td[:k]:
    if ti in kindset:
        heappush(hque, (di, ti))
        
    else:
        kindnum += 1
        kindset.add(ti)
        
    ans += di
    
ans += len(kindset)*len(kindset)
    
rems = deque([])
for ti, di in td[k:]:
    if ti in kindset:
        continue
    
    else:
        rems.append((di, ti))


score = ans

for _ in range(min(len(hque), len(rems))):
    old_d, _ = heappop(hque)
    new_d, new_t = rems.popleft()
    
    if new_t in kindset:
        continue
    
    else:
        score = score + new_d - old_d
        score += (2*kindnum + 1)
        
        kindnum += 1
        kindset.add(new_t)
        
        ans = max(ans, score)
        
print(ans)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

def inpl(): return [int(i) for i in input().split()]

N,K = inpl()
d = {}
ans = 0
n0 = []
n1 = []

for i in range(N):
    ti,di = inpl()
    if (ti in d.keys()):
        d[ti].append(di)
    else:
        d[ti] = [di]

for i in d.keys():
    sortedd = sorted(d[i])
    sortedd.reverse()
    n1.append(d[i][0])
    n0 += d[i][1:]

n0.sort()
n0.reverse()
n1.sort()
n1.reverse()

for i in range(len(n1)+1):
    if(K-i <= len(n0) and 0<=K-i):
        total = sum(n1[:i])+sum(n0[:K-i])+(i)**2
        if(ans<total):
            ans = total

print(ans)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

def inpl(): return [int(i) for i in input().split()]

N,K = inpl()
td = [0]*N
d = {}
ans = 0
for i in range(N):
    td[i] = inpl()
    if (td[i][0] in d.keys()):
        d[td[i][0]].append(td[i][1])
    else:
        d[td[i][0]] = [td[i][1]]

p = len(d)
n0 = []
n1 = []

for i in d.keys():
    sortedd = sorted(d[i])
    sortedd.reverse()
    for index,value in enumerate(sortedd):
        if(index == 0):
            n1.append(value)
        else:
            n0.append(value)

n0.sort()
n0.reverse()
n1.sort()
n1.reverse()

for i in range(len(n1)+1):
    if(K-i <= len(n0)):
        total = sum(n1[:i])+sum(n0[:K-i])+(i)**2
        if(ans<total):
            ans = total

print(ans)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k != 0:
        if a[A[i][0]] ==0:
            a[A[i][0]] = 1
            eat.append(A[i][1])
            k -= 1
        else:
            noteat.append(A[i][1])
    else:
        noteat.append(A[i][1])
        
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat[k:])+1)):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat[k:])+1)):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Submission status: Wrong Answer

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

inputs = {}
priority_queue = []
queue = []
ans = 0

for _ in range(N):
    key, value = map(int, input().split())
    if key not in inputs:
        inputs[key] = [value]
    else:
        inputs[key].append(value)

for key in inputs:
    inputs[key].sort(reverse=True)
    priority_queue.append(inputs[key][0])
    queue += inputs[key][1:]

priority_queue.sort(reverse=True)
queue.sort(reverse=True)

for i in range(1, len(priority_queue)+1):
    total = sum(priority_queue[:i]) + sum(queue[:K-i]) + i ** 2
    ans = max(ans, total)

print(ans)

Submission status: Wrong Answer

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

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

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

sushi_list.sort(key=lambda tup: tup[1])

max_sushi_list, rest_of_sushi = sushi_list[-k:], sushi_list[:-k]
max_neta_set = {sushi[0] for sushi in max_sushi_list}

iMax = len(neta_set)
iMin = len(max_neta_set)

max_score = iMin**2 + sum([sushi[1] for sushi in max_sushi_list])


for nNeta in range(iMin+1, iMax + 1, 1):

    # Find sushi having minimum oishisa and being not unique for neta
    r = len(max_sushi_list)
    swap_index = -1
    for i in range(r):
        if max_sushi_list[i][0] in [neta for neta,oishisa in max_sushi_list[:i]+max_sushi_list[i+1:]]:
            swap_index = i
            break

    # Find sushi having maximum oishisa and unique for neta
    r = len(rest_of_sushi)
    for i in range(r):
        if not rest_of_sushi[r-i-1][0] in max_neta_set:
            rest_of_sushi.append(max_sushi_list.pop(swap_index))
            max_sushi_list.append(rest_of_sushi.pop(r-i-1))
            break

    score = nNeta**2 + sum([sushi[1] for sushi in max_sushi_list])
    if score > max_score:
      max_score = score
    sushi_list.sort(key=lambda tup: tup[1])
    max_neta_set = {sushi[0] for sushi in max_sushi_list}

print(max_score)

Submission status: Wrong Answer

Solution code:
N,K= map(int,input().split())
n=0
lis=[]
points=0
kind=set()
l1=[]
l2=[]
for _ in range(N):
    t,d = map(int,input().split())
    lis.append((t,d))
lis.sort(key=lambda x:x[1],reverse=True)
for i in lis:
    if K>n:
        if i[0] not in kind:
            l1.append(i[1])
            kind.add(i[0])
        else:
            l2.append(i[1])
        n+=1
        if K==n:
            v=len(kind)
            tl=sum(l1) +sum(l2) +v**2
            points=max(tl,points)
            if len(l2)==0:
                break
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i[1])
            v+=1
            tl=sum(l1) +sum(l2) +v**2
            points=max(tl,points)
            if len(l2)==0:
                break
print(points)

Submission status: Wrong Answer

Solution code:
N,K= map(int,input().split())
n=0
lis=[]
points=0
kind=set()
kindall=set()
l1=[]
l2=[]
for _ in range(N):
    t,d = map(int,input().split())
    kindall.add(t)
    lis.append((t,d))
lis.sort(key=lambda x:x[1],reverse=True)
up=min(K,len(kindall))
for i in lis:
    if K>n:
        if i[0] not in kind:
            l1.append(i[1])
            kind.add(i[0])
        else:
            l2.append(i[1])
        n+=1
        if K==n:
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i[1])
            kind.add(i[0])
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
            
    if len(kind)==up:
        break

print(points)

Submission status: Wrong Answer

Solution code:
N,K= map(int,input().split())
n=0
lis=[]
points=0
kind=set()
kindall=set()
l1=[]
l2=[]
for _ in range(N):
    t,d = map(int,input().split())
    kindall.add(t)
    lis.append((t,d))
lis.sort(key=lambda x:x[1],reverse=True)
up=min(K,len(kindall))
for i in lis:
    if K>n:
        if i[0] not in kind:
            l1.append(i)
            kind.add(i[0])
        else:
            l2.append(i)
        n+=1
        if K==n:
            tl=0
            for k in l1:
                tl+=k[1]
            for k in l2:
                tl+=k[1]
            tl+=len(kind)**2
            points=max(tl,points)
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i)
            kind.add(i[0])
            tl=0
            for k in l1:
                tl+=k[1]
            for k in l2:
                tl+=k[1]
            tl+=len(kind)**2
            points=max(tl,points)
            
    if len(kind)==up:
        break

print(points)

Submission status: Wrong Answer

Solution code:
def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

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

L_sorted=sorted(L, key=lambda x:x[1], reverse=True)

L_kind=set([i[0] for i in L])
L_sorted_kind=set([i[0] for i in L_sorted[:K]])

tmp=L_sorted[K:]
L_sorted_=[]
tmp_set=[]
for i in tmp:
  if not (i[0] in set(L_sorted_kind)) and not (i[0] in set(tmp_set)):
    tmp_set.append(i[0])
    L_sorted_.append(i)

max_score=score(L_sorted[:K]) 
  
print(max_score)

Submission status: Wrong Answer

Solution code:
N,K=map(int,input().split())
b=[] #[[種類,美味しさ],...]
 
for _ in range(N):
    t,d=map(int,input().split())
    b.append([t,d])
b.sort(key=lambda x:-x[1])
 
f=[0]*(N+1)
x=0
point=0
 
for [t,d] in b[:K]:
    if f[t]==0: x+=1
    f[t]+=1
    point+=d
maxpoint=point+x*x
 
h=K-1
for [t,d] in b[K:]:
    if f[t]==0:
        f[t]=1
        x+=1
        point+=d
        for i in range(h,-2,-1):
            if i==-1:
                print(maxpoint)
                exit()
            t=b[i][0]
            d=b[i][1]
            if f[t]>1:
                f[t]-=1
                point-=d
                break
        maxpoint=max(maxpoint,point+x*x)
print(maxpoint)

Submission status: Wrong Answer

Solution code:
def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

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

L_sorted=sorted(L, key=lambda x:x[1], reverse=True)

L_kind=set([i[0] for i in L])
L_sorted_kind=set([i[0] for i in L_sorted[:K]])

bucket={k:len([k for i in L_sorted[:K] if i[0]==k]) for k in L_sorted_kind}

tmp=L_sorted[K:]
L_sorted_=[]
tmp_set=[]
for i in tmp:
  if not (i[0] in set(L_sorted_kind)) and not (i[0] in set(tmp_set)):
    tmp_set.append(i[0])
    L_sorted_.append(i)

max_score=score(L_sorted[:K])
  
print(max_score)

Submission status: Wrong Answer

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

def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2


L=[list(map(int, input().split())) for _ in range(N)]
L_sorted=sorted(L, key=lambda x:x[1], reverse=True)

L_kind=set([i[0] for i in L])
L_sorted_kind=set([i[0] for i in L_sorted[:K]])

bucket={k:len([k for i in L_sorted[:K] if i[0]==k]) for k in L_sorted_kind}

tmp=L_sorted[K:]
L_sorted_=[i for i in tmp if len(set([i[0]]) & set(L_sorted_kind))==0]

max_score=score(L_sorted[:K])

i=1
while len(L_kind & L_sorted_kind)!=len(L_kind) and i<=K:
  if bucket[L_sorted[:K][-i][0]]>1:
    bucket[L_sorted[:K][-i][0]]-=1
    
    L_sorted[K-i]=L_sorted_[0]
    L_sorted_.pop(0)
    
    L_sorted_kind=set([i[0] for i in L_sorted[:K]])

  tmp_score=score(L_sorted[:K])
  if tmp_score>max_score:
    max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Wrong Answer

Solution code:
import numpy as np

m = list(map(int, input().split()))
n = m[0]
k = m[1]
a = np.zeros((0,2))
score = k * k
k1 = k
for i in range(n):
    b = np.asarray(list(map(int, input().split())))
    a = np.r_[a,b.reshape(1,-1)]

a = np.insert(a,2,0,axis=1)
#print(a) # [a1, a2, a3, ..., aN]
for _ in range(k):
    am = np.argmax(a[:,1])
    tmp=a[am,0]
    score += a[am,1]
    if a[am,2] == 1:
        k1 -= 1
    if a[am,2] == 0:
        a[np.where(a[:,0] == tmp)[0],1] -= 2*k1 - 1
        a[np.where(a[:,0] == tmp)[0],2] = 1
        #print(a)
    

    
    a[am]= [0,0,1]
    #print(tmp)

    #print(a[np.where(a[:,0] == tmp)[0],1]) 

print(int(score))

Submission status: Wrong Answer

Solution code:
import numpy as np

m = list(map(int, input().split()))
n = m[0]
k = m[1]
a = np.zeros((0,2))
score = k * k
k1 = k
for i in range(n):
    b = np.asarray(list(map(int, input().split())))
    a = np.r_[a,b.reshape(1,-1)]

a = np.insert(a,2,0,axis=1)
#print(a) # [a1, a2, a3, ..., aN]
for _ in range(k):
    am = np.argmax(a[:,1])
    tmp=a[am,0]
    score += a[am,1]
    #print(a[am])
    a[am]= [0,0,1]
    #print(tmp)
    a[np.where(a[:,0] == tmp)[0],1] -= 2*k1 - 1
    k1 -= 1
    #print(a[np.where(a[:,0] == tmp)[0],1]) 

print(int(score))

Submission status: Wrong Answer

Solution code:
import numpy as np

m = list(map(int, input().split()))
n = m[0]
k = m[1]
a = np.zeros((0,2))
score = k * k
k1 = k
for i in range(n):
    b = np.asarray(list(map(int, input().split())))
    a = np.r_[a,b.reshape(1,-1)]

a = np.insert(a,2,0,axis=1)
print(a) # [a1, a2, a3, ..., aN]
for _ in range(k):
    am = np.argmax(a[:,1])
    tmp=a[am,0]
    score += a[am,1]
    #print(a[am])
    a[am]= [0,0,1]
    print(tmp)
    a[np.where(a[:,0] == tmp)[0],1] -= 2*k1 - 1
    k1 -= 1
    print(a[np.where(a[:,0] == tmp)[0],1]) 

print(score)

Submission status: Wrong Answer

Solution code:
n,k = map(int, input().split())
td = []
shurui_list = []
for _ in range(n):
    tdi = list(map(int, input().split()))
    td.append(tdi)
    shurui_list.append(tdi[0])
sorted(td, key=lambda x: -x[1])
shrui_list = set(shurui_list)

oisisa = 0
shurui = []
max1 = 0
i = 0

oisisa2 = 0
shurui2 = []
hoketsu = []
max2 = 0
j = 0
for temp in td:
    # ネタ
    if i == k:
        break
    oisisa += temp[1]
    shurui.append(temp[0])
    i += 1

    # 種類
    if j == k:
        break
    if temp[0] in shurui2 and len(shurui2) < len(shrui_list):
        hoketsu.append(temp[1])
        continue
    if len(shurui2) == len(shrui_list):
        for h in hoketsu:
            oisisa2 += h
            j += 1
            if j == k:
                break
    oisisa2 += temp[1]
    shurui2.append(temp[0])
    j += 1

max1 = len(set(shurui))**2 + oisisa
max2 = len(set(shurui))**2 + oisisa2

# 比較
max = max1 if max1 > max2 else max2
print(max)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter
import heapq
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]
td.sort(key=itemgetter(1, 0), reverse=True)
t, d = zip(*td)

s = set()
q = []
#heapq.heapify(q)
#q = deque(q)
tmp = 0
now = 0
for i in range(K):
    if t[i] not in s:
        s.add(t[i])
    else:
        #heapq.heappush(q, d[i])
        q.append(d[i])

    tmp += d[i]
    now += 1
kind = len(s)
sums = [-100000000000]*(kind+1)
sums[kind] = tmp
while q:
    kind += 1
    #tmp = -heapq.heappop(q)
    tmp = q.pop()
    for i in range(now, N):
        now += 1
        if t[i] not in s:
            sums.append(sums[kind-1] + d[i] - tmp)
            break
for i in range(len(sums)):
    sums[i] += i**2
print(max(sums))

Submission status: Wrong Answer

Solution code:
from operator import itemgetter
from collections import Counter
import heapq
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]
td.sort(key=itemgetter(1, 0), reverse=True)
t, d = zip(*td)

s = set()
q = []
heapq.heapify(q)
tmp = 0
now = 0
for i in range(K):
    if t[i] not in s:
        s.add(t[i])
    else:
        heapq.heappush(q, d[i])
    tmp += d[i]
    now += 1
kind = len(s)
sums = [-100000000000]*(kind+1)
sums[kind] = tmp
while q:
    kind += 1
    tmp = -heapq.heappop(q)
    for i in range(now, N):
        now += 1
        if t[i] not in s:
            tmp += d[i]
            sums.append(sums[kind-1] + tmp)
            break
for i in range(len(sums)):
    sums[i] += i**2
print(max(sums))

Submission status: Wrong Answer

Solution code:
from collections import  deque
N,K = map(int, input().split())
neta_uma = []
for i in range(N):
    neta,uma = map(int, input().split())
    neta_uma.append((neta, uma))
neta_uma = sorted(neta_uma, key=lambda x:x[1], reverse=True)

# うまさ全振り
chosen = neta_uma[:K]
used = set(n for n,u in chosen)
# print(chosen, used)
ans = sum(u for n,u in chosen) + len(used)**2

# 多様性を増やす方向
tmp = ans
for i in range(K,N):
    # print(neta_uma[i], ans, used, )
    if not chosen: break
    
    if neta_uma[i][0] not in used:
        n,u = chosen.pop()
        tmp += neta_uma[i][1] - u
        tmp += 2*len(used) +1     # (k+1)^2 - k^2
        used.add(neta_uma[i][0])
        ans = max(ans, tmp)
        
print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
sushi_list = []
max_d = [0] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi_list.append((t - 1, d))
    max_d[t - 1] = max(max_d[t - 1], d)

sushi_list.sort(key=lambda x: x[1], reverse=True)
good = []
for i in range(N):
    if max_d[sushi_list[i][0]] == sushi_list[i][1]:
        good.append(True)
    else:
        good.append(False)

s = sum([sushi[1] for sushi in sushi_list[:K]])
p = len(set([sushi[0] for sushi in sushi_list[:K]]))
ans = s + p**2
r = K - 1
l = K
while True:
    while r >= 0 and good[r]:
        r -= 1
    while l < N and not good[l]:
        l += 1
    if r < 0 or l >= N:
        break
    s = s - sushi_list[r][1] + sushi_list[l][1]
    p += 1
    ans = max(ans, s + p**2)
    r -= 1
    l += 1

print(ans)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter
from collections import Counter
import heapq
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]
td.sort(key=itemgetter(1, 0), reverse=True)
t, d = zip(*td)

s = set()
q = []
heapq.heapify(q)
tmp = 0
for i in range(K):
    if t[i] not in s:
        s.add(t[i])
    else:
        heapq.heappush(q, d[i])
    tmp += d[i]
kind = len(s)
sums = [0]*(kind+1)
sums[kind] = tmp
while q:
    kind += 1
    tmp = -heapq.heappop(q)
    for i in range(N):
        if t[i] not in s:
            tmp += d[i]
            sums.append(sums[kind-1] + tmp)
            break
for i in range(len(sums)):
    sums[i] += i**2
print(max(sums))

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
from collections import deque
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])

used = set()
q = []
res = 0
kind = 0
for i in range(k):
    t,d = td[i]
    if t not in used:
        kind += 1
        used.add(t)
    else:
        q.append((t,d))
    res += d
res += kind**2
q = deque(sorted(q, key=lambda x:x[1]))

tmp = res
for i in range(k,n):
    if len(q)==0:
        break
    t2,d2 = td[i]
    if t2 in used:
        continue
    t,d = q.popleft()
    used.add(t)

    r = tmp - d + d2
    r -= kind**2
    kind += 1
    r += kind**2
    res = max(res, r)
    tmp = r

print(res)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
from collections import deque
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])

used = set()
q = []
res = 0
kind = 0
for i in range(k):
    t,d = td[i]
    if t not in used:
        kind += 1
        used.add(t)
    else:
        q.append((t,d))
    res += d
res += kind**2
q = deque(sorted(q, key=lambda x:x[1]))
if len(q)==0:
    print(res)
    exit()

tmp = res
for i in range(k,n):
    t2,d2 = td[i]
    if t2 in used:
        continue
    t,d = q.popleft()

    r = tmp - d + d2
    r -= kind**2
    kind += 1
    r += kind**2
    res = max(res, r)
    tmp = r
    if len(q)==0:
        break

print(res)