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:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    if count < K:
        while 1:
            if count >= K :
                break
            if len(not_eat_list) <= 0:
                break
            i =not_eat_list.popleft()
            eat_list.appendleft(i)
            max_score += sushi[i][0]
        print max_score
        return


    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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

s = []
for i in range(N):
  s.append(list(map(int, input().split())))

s2 = sorted(s ,key=lambda x: x[1], reverse=True)
s3 = sorted(s2,key=lambda x: x[0])

res = s2[:K]
p1 = sum([i[1] for i in res]) + len(set([i[0] for i in res]))**2

n1 = set([i[0] for i in s2])

res2 = [s3[0]]
res3 = []

for i in range(1,N):
    if s[i-1][0] != s[i][0]:
        res2.append(s[i])
    else:
        res3.append(s[i])

res2 = sorted(res2 ,key=lambda x: x[1], reverse=True)
res3 = sorted(res3 ,key=lambda x: x[1], reverse=True)

if len(res2) >= K:
    res2 = res2[:K]
    p2 = sum([i[1] for i in res2]) + len(set([i[0] for i in res2]))**2
elif K - len(res2) == 1:
    res2.append(res3[0])
    p2 = sum([i[1] for i in res2]) + len(set([i[0] for i in res2]))**2
else:
    res2.expend(res3[:K-len(res2)])
    p2 = sum([i[1] for i in res2]) + len(set([i[0] for i in res2]))**2
    


print(max(p1,p2))

Submission status: Wrong Answer

Solution code:
import numpy as np

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

sushi = np.array( [list(map(int, input().split())) for i in range(n)] )
index=list(range(0,n))

ans=0
for i in range(10000):
    tmp = np.random.choice(index, k, replace=False)
    sushi_ = sushi[tmp]
    ans_   = np.sum(sushi_[:,1]) + len(set(sushi_[:,0]))**2
    if ans_ > ans : ans = ans_

print(ans)

Submission status: Wrong Answer

Solution code:
import numpy as np

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

sushi = np.array( [list(map(int, input().split())) for i in range(n)] )
index=list(range(0,n))

ans=0
for i in range(1000):
    tmp = np.random.choice(index, k, replace=False)
    sushi_ = sushi[tmp]
    ans_   = np.sum(sushi_[:,1]) + len(set(sushi_[:,0]))**2
    if ans_ > ans : ans = ans_

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(already_eat)

    while 1:
        if count >= K :
            break
        if len(not_eat_list) <= 0:
            break
        i =not_eat_list.popleft()
        eat_list.appendleft(i)
        max_score += sushi[i][0]


    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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

ate=[0 for i in range(N+1)]
sushi=[]

types=set()

for i in range(N):
    sushi.append(list(map(int,input().split())))
    types.add(sushi[i][0])
types=len(types)


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

x=0
res=0

for s in sushi[:K]:
    if ate[s[0]]==0:
        x+=1
    ate[s[0]]+=1
    res+=s[1]

repl=K

memo=[res+x**2]

rem=[0 for i in range(K)]

while x<types and repl<=N-1:
    for i in range(K):
        t,d=sushi[K-i-1]
        if ate[t]>1 and rem[K-i-1]==0:
            ate[t]-=1
            rem[K-i-1]+=1
            res-=d
    flag=True
    while repl<=N-1:
        t,d=sushi[repl]
        if ate[t]==0:
            ate[t]+=1
            x+=1
            res+=d
            break
        repl+=1
    memo.append(res+x**2)


print(max(memo))

Submission status: Wrong Answer

Solution code:
import sys
from collections import deque
import copy
import math
def get_read_func(fileobject):
    if fileobject == None :
        return raw_input
    else:
        return fileobject.readline

def main():
    if len(sys.argv) > 1:
        f = open(sys.argv[1])
    else:
        f = None
    read_func = get_read_func(f);
    input_raw = read_func().strip().split()
    [N, K] = [long(input_raw[0]), long(input_raw[1])]
    sushi = []
    for i in range(N):
        input_raw = read_func().strip().split()
        [t, d] = [long(input_raw[0]), long(input_raw[1])]
        sushi.append((d, t))

    sushi.sort(reverse=True)

    already_eat = set([])
    max_score = 0L
    count = 0
    not_eat_list = deque([])
    eat_list = deque([])
    for i in range(N):
        if sushi[i][1] not in already_eat:
            already_eat.add(sushi[i][1])
            max_score += sushi[i][0]
            eat_list.append(i)
        else:
            not_eat_list.append(i)

        if len(already_eat) == K:
            break
    count = len(eat_list)

    while 1:
        if count >= K :
            break
        if len(not_eat_list) <= 0:
            break
        i =not_eat_list.popleft()
        eat_list.appendleft(i)
        max_score += sushi[i][0]




    now_kind = len(already_eat)
    max_score += now_kind * now_kind
    now_score = max_score
    for new_idx in not_eat_list:
        if len(eat_list) == 0 :
            break
        eated_idx = eat_list.pop()
        already_eat.remove(sushi[eated_idx][1])
        already_eat.add(sushi[new_idx][1])

        new_kind = len(already_eat)
        now_score =  now_score - (now_kind * now_kind) + (new_kind *new_kind) - sushi[eated_idx][0] + sushi[new_idx][0]
        now_kind = new_kind
        max_score =max(max_score, now_score)


    print max_score




if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1], reverse=True)
#print(sushi)
num = 0
kind = set()
changable = []
for i in range(k):
	num += sushi[i][1]
	if sushi[i][0] in kind:
		changable.append(sushi[i][1])
	else:
		kind.add(sushi[i][0])

num += len(kind)**2
changable.reverse()
#print(changable)
ans = num

for i in range(k, n):
	if not changable:
		break

	if sushi[i][0] not in kind:
		num += sushi[i][1] + 2*len(kind)+1 - changable.pop()
		#print(num, kind)
		kind.add(sushi[i][0])
		ans = max(ans, num)


print(ans)

Submission status: Wrong Answer

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

sushi.sort()
temp = sushi[0][0]
r = 0
sus = [[]]
for s in sushi:
  if s[0] == temp:
    sus[r] = sus[r] + [s[1]]
  else:
    sus[r].sort()
    sus[r].reverse()
    r += 1
    temp = s[0]
    sus = sus + [[s[1]]]
sus[r].sort()
sus[r].reverse()
sus.sort()
sus.reverse()

x = 0
q = 0
u = []
v = []
for i in range(len(sus)):
  q += len(sus[i])
  u = u + sus[i][1:]
  v.append(sus[i][0])
  if q >= k:
    v = [sus[j][0] for j in range(i+1)]
    u.sort()
    u.reverse()
    w = sum(u[:k-(i+1)])
    x = max(x, (i + 1) ** 2 + sum(v) + w)
print(x)

Submission status: Wrong Answer

Solution code:
if __name__ == '__main__':
    import sys
    input = sys.stdin.readline
    import collections
    n, k = map(int, input().split())
    td = [list(map(int, input().split())) for _ in range(n)]
    td.sort(key=lambda x: x[1], reverse=True)
    ans_list = [0]
    reverse_td = list(zip(*td))
    c = collections.Counter(reverse_td[0][:k])
    kind = len(c)
    ans = sum(reverse_td[1][:k])+kind**2
    check = k-1
    flag=False
    for i in range(k,n):
        t = td[i][0]
        d = td[i][1]
        if t in c:
            pass
        else:
            c[t]=1
            while True:
                new_t = td[check][0]
                if c[new_t]==1:
                    check -=1
                    if check == -1:
                        flag =True
                        break
                else:
                    c[new_t]-=1
                    ans_list.append(2*kind+1+d-td[check][1])
                    check-=1
                    kind+=1
                    break
            if flag:
                break
    print(ans+max(ans_list))

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
td = [list(map(int, input().split())) for _ in range(n)]
td.sort(key=lambda x:x[1],reverse=True)
a = 0
b = 0
v1 = []
v2 = []
cnt = 0
for i in range(n):
    if i < k:
        a += td[i][1]
        if not td[i][0] in v1:
            v1.append(td[i][0])
    if not td[i][0] in v2:
        v2.append(td[i][0])
        b += td[i][1]
        cnt += 1
    if cnt == k:
        break
a += len(v1)**2
b += len(v2)**2
print(max(a,b))

Submission status: Wrong Answer

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

sushi.sort()
temp = sushi[0][0]
r = 0
sus = [[]]
for s in sushi:
  if s[0] == temp:
    sus[r] = sus[r] + [s[1]]
  else:
    sus[r].sort()
    sus[r].reverse()
    r += 1
    temp = s[0]
    sus = sus + [[s[1]]]
sus[r].sort()
sus[r].reverse()
sus.sort()
sus.reverse()
print(sus)

x = 0
q = 0
u = []
v = []
for i in range(len(sus)):
  q += len(sus[i])
  u = u + sus[i][1:]
  v.append(sus[i][0])
  if q >= k:
    v = [sus[j][0] for j in range(i+1)]
    u.sort()
    u.reverse()
    w = sum(u[:k-(i+1)])
    x = max(x, (i + 1) ** 2 + sum(v) + w)
print(x)

Submission status: Wrong Answer

Solution code:
from statistics import mean, median,variance,stdev
import sys
import math
import fractions

def j(q):
    if q==1: print("YES")
    else:print("NO")
    exit(0)


def ct(x,y):
    if (x>y):print("")
    elif (x<y): print("")
    else: print("")

def ip():
    return int(input())

#x = ip()                              #入力整数1つ
n,k= (int(i) for i in input().split())      #入力整数横2つ
#x,y,z = (int(i) for i in input().split())    #入力整数横3つ
#n,m,x,y= (int(i) for i in input().split())  #入力整数横4つ
#a = [int(i) for i in input().split()]        #入力整数配列
#a = input()                                  #入力文字列
#a = input().split()                          #入力文字配列

#jの変数はしようできないので注意
a = []
o = []
p = []
for i in range(n):
    c,d = (int(i) for i in input().split())
    a.append([c,d])
    o.append([d,c])
    p.append(d)
o.sort(reverse=True)
#print(o)
l = a[0][0]
b = []
c = []

maxl = []

for i in range(len(a)):
#    print(a[i][1],l)
    if l == a[i][0]: b.append(a[i][1])
    else :
        c.append(b)
        b = []
        b.append(a[i][1])
        l = a[i][0]
c.append(b)
neta  = 0
deli = 0
#print(c.sort(reverse=True),len(c))

for i in range(len(c)):
#    print(i)
    a = c[i].pop(0)
    neta+=1
    deli+=a
    if k==i+1:break
b = []
i+=1
u = k
k-=i
#print(c)
if k:
    for i in range(len(c)):
        b.extend(c[i])
    b.sort(reverse=True)
    for i in range(k):
        deli+=b[i]
#print(deli,neta)
temp = deli+neta*neta
d = []
deli = 0
neta = 0
#print(u,o)

for i in range(u):
    if o[i][1] not in d:
        d.append(o[i][1])
        neta+=1
 #       print(d)
    deli+=o[i][0]
 #   print(o)
#print(deli,neta)
if deli+neta*neta > temp:print(deli+neta*neta)
else: print(temp)

Submission status: Wrong Answer

Solution code:
import math
import copy




def nCast(number):
    if type(number)==str:
        return int(number)
    for idx in range(0,len(number)):
        if type(number[idx])==str:

            number[idx]=int(number[idx])
        else:
            nCast(number[idx])
    return number
    
def inputArr(w):
    l=list()
    for idx in range(0,w):
        l.append(input())
    return l
def inputArr1(w):
    l=list()
    for idx in range(0,w):
        l.append(input().split())
    return l

n=nCast(input().split())
k=nCast(inputArr1(n[0]))
def sortKey(k):
    return k[1]
k.sort(key=sortKey,reverse=True)

a=[0,1,2]

def func():
    d=dict()
    result=0
    s=set()
    for idx in range(0,len(k)):
        s.add(k[idx][0])
    for idx in range(0,n[1]):
        if k[idx][0] not in d:
            d[k[idx][0]]=0
        d[k[idx][0]]+=1
        result+=k[idx][1]
    tmp=result+len(d)*len(d)
    up=n[1]-1
    down=n[1]
    while up>=0 and down<len(k):

        while up>=0:
            if d[k[up][0]]>1:
                result-=k[up][1]
                d[k[up][0]]-=1
                up-=1
                break
            else:
                up-=1   
        while down<len(k):
            if k[down][0] not in d:
                d[k[down][0]]=1
                result+=k[down][1]
                tmp=max(tmp,result+len(d)*len(d))
                down +=1
                break
            else:
                down+=1
    return tmp
print(func())

Submission status: Wrong Answer

Solution code:
import math
import itertools

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

dt = []
neta = []
bass_score = 0
netamax = []
for i in range(N):
    if i == 0:
        tmp1,tmp2 = map(int, input().split())
        netamax.append(tmp2)
        dt.append([tmp2,tmp1])
    else:
        tmp1,tmp2 = map(int, input().split())
        if dt[i-1][1] != tmp1:
            netamax.append(tmp2)
        dt.append([tmp2,tmp1])


DT = dt.copy()
DT.sort(reverse=True)

for i in range(K):
    bass_score += DT[i][0]
    neta.append(DT[i][1])

bonus_score = len(set(neta))**2
test1 = bass_score + bonus_score
#----------------------------
netamax.sort(reverse=True)
bass_score = sum(netamax[:K])

test2 = bass_score + K*K
print(max(test1,test2))

Submission status: Wrong Answer

Solution code:
#!usr/bin/env python3
def LI(): return list(map(int, input().split()))
def II(): return int(input())
def LS(): return input().split()
def S(): return input()
def LIR(n): return [LI() for i in range(n)]
mod = 1000000007

#A
"""
a,b,c = LI()
s = a+b+c
s /= 2
ans = (s*(s-a)*(s-b)*(s-c))**(1/2)
print(int(ans))
"""

#B
"""
s = int(input())
m = 1
lis = [True for i in range(1000001)]

while lis[s]:
    lis[s] = False
    if s % 2 == 1:
        s = 3*s+1
    else:
        s //= 2
    m += 1
print(m)
"""

#C
"""
n = II()
h = LI()
ans = 0
while sum(h) != 0:
    l = 0
    while h[l] == 0: l += 1
    r = l
    while r < n:
        if h[r] == 0: break
        h[r] -= 1
        r += 1
    ans += 1
print(ans)
"""

#D
n,k = LI()
lis = LIR(n)
lis.sort(key = lambda x:x[1])
lis = lis[::-1]
num = [False for i in range(n)]
count = [0 for i in range(n)]
for i in range(n):
    lis[i][0] -= 1
ans = 0
t = 0
l = []
for i in range(k):
    ans += lis[i][1]
    count[lis[i][0]] += 1
    if not num[lis[i][0]]:
        t += 1
    num[lis[i][0]] = True
for i in range(k):
    if count[i] != 1:
        l.append(lis[i])
l = l[::-1]

a = [0 for i in range(len(l))]
for i in range(len(l)):
    a[i] = l[i][1]
for i in range(len(l)-1):
    a[i+1] += a[i]
li = []

for i in range(n):
    if not num[lis[i][0]]:
        li.append(lis[i][1])
        num[lis[i][0]] = True
b = [0 for i in range(len(li))]
for i in range(len(li)):
    b[i] = li[i]
for i in range(len(li)-1):
    b[i+1] += b[i]
d = [[ans,t] for i in range(min(len(l),len(li))+1)]

for i in range(len(d)-1):
    d[i+1][0] += b[i] - a[i]
    d[i+1][1] = min(k,t+i+1)
ans = 0
for i in range(len(d)):
    ans = max(ans, d[i][0]+d[i][1]**2)
print(ans)
#E

#F

#G

#H

#I

#J

#K

#L

#M

#N

#O

#P

#Q

#R

#S

#T

Submission status: Wrong Answer

Solution code:
#!/Users/takahiro.kida/git/atcorder/venv/bin/python

import numpy
import scipy

from collections import defaultdict
from pprint import pprint

N, K = [int(x) for x in input().split(" ")]
sushis = [] # type, delicious, eat?
for n in range(N):
    row = [int(x) for x in input().split(" ")] + [False]
    sushis.append(row)



type_of_max = {}  # type: [index of sushis, deli]
for i in range(N):
    sushi = sushis[i]
    if not sushi[0] in type_of_max:
        type_of_max[sushi[0]] = [i, sushi[1]]
        continue
    if type_of_max[sushi[0]][1] < sushi[1]:
        type_of_max[sushi[0]] = [i, sushi[1]]
        continue


vali = 0
eat = []
for type in type_of_max:
    index = type_of_max[type][0]
    eat.append(sushis[index])
    sushis[index][2] = True
    vali += 1

eat = sorted(eat, key=lambda sushi: sushi[1], reverse=True)
dish = []
for sushi in sushis:
    if not sushi[2]:
        dish.append(sushi)
dish = sorted(dish, key=lambda sushi: sushi[1], reverse=True)

if len(eat) > K:
    for i in range(K, len(eat) - 1):
        sushi = eat.pop()
        dish.append(sushi)
        vali += -1

if len(eat) < K:
    for i in range(0, K - len(eat) -1):
        sushi = dish.pop()
        eat.append(sushi)

dish = sorted(dish, key=lambda sushi: sushi[1], reverse=True)
eat = sorted(eat, key=lambda sushi: sushi[1], reverse=True)

while True:
    if not dish:
        break
    d_diff = dish[0][1] - eat[-1][1]
    t_diff = vali * vali - (vali - 1) * (vali - 1)
    if (d_diff > t_diff):
        eat.pop()
        eat.insert(0, dish.pop(0))
        vali += -1
    else:
        break

deli = [x[1] for x in eat]
print(vali*vali + sum(deli))

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

n, k = map(int, input().split(" "))
sushis = []
flags = [False for i in range(n+10)]
sushis2 = []
sushis22 = []
sum2 = 0
count2 = 0
for i in range(n):
    t, d = (map(int, input().split(" ")))
    sushis.append([t, d])

    if not flags[t]:
        sushis22.append(d)
        flags[t] = True
    else:
        sushis2.append(d)
# calc variety
if len(sushis22) <= k:
    count2 += len(sushis22)
    sum2 += sum(sushis22)
    sum2 += len(sushis22) * len(sushis22)
else:
    sum2 += sum(sorted(sushis22, reverse=True)[0:k])
    sum2 += k*k
    count2 += k

if count2 < k:
    sushis2 = sorted(sushis2, reverse=True)
    sum2 += sum(sushis2[0:k-count2])

# calc score
sum1 = 0
set1 = []
count = k
temp_score = 0
temp_i = 0
sushis = sorted(sushis, key=itemgetter(1), reverse=True)

temp_set = []
for i, sushi in enumerate(sushis):
    sum1 += sushi[1]
    if count == 1:
        temp_set.append(sushi[0])
        temp_i = i
        temp_score = sushi[1]
        break
    group = sushi[0]
    set1.append(group)
    count -= 1


set1 = set(set1)

while temp_i+1 < n:
    if temp_score == sushis[temp_i+1]:
        temp_set.append(sushis[temp_i+1][0])
        temp_i += 1
    else:
        break

if len(set1 | set(temp_set)) != len(set1):
    sum1 += (len(set1)+1)*(len(set1)+1)
else:
    sum1 += len(set1) * len(set1)
print(max(sum1, sum2))

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

n, k = map(int, input().split(" "))
sushis = []
flags = [False for i in range(n+10)]
sushis2 = []
sushis22 = []
sum2 = 0
count2 = 0
for i in range(n):
    t, d = (map(int, input().split(" ")))
    sushis.append([t, d])

    if not flags[t]:
        sushis22.append(d)
        flags[t] = True
    else:
        sushis2.append(d)
# calc variety
if len(sushis22) <= k:
    count2 += len(sushis22)
    sum2 += sum(sushis22)
    sum2 += len(sushis22) * len(sushis22)
else:
    sum2 += sum(sorted(sushis22, reverse=True)[0:k])
    sum2 += k*k
    count2 += k

if count2 < k:
    sushis2 = sorted(sushis2, reverse=True)
    sum2 += sum(sushis2[0:k-count2])

# calc score
sum1 = 0
set1 = []
count = k
temp_score = 0
temp_i = 0
sushis = sorted(sushis, key=itemgetter(1), reverse=True)

temp_set = []
for i, sushi in enumerate(sushis):
    sum1 += sushi[1]
    if count == 1:
        temp_set.append(sushi[0])
        temp_i = i
        temp_score = sushi[1]
        break
    group = sushi[0]
    set1.append(group)
    count -= 1


set1 = set(set1)

while temp_i+1 < n:
    if temp_score == sushis[temp_i+1]:
        temp_set.append(sushis[temp_i+1][0])
        temp_i += 1
    else:
        break

if len(set1 | set(temp_set)) != len(set1):
    sum1 += (len(set1)+1)*(len(set1)+1)
else:
    sum1 += len(set1) * len(set1)
print(sum1)
print(sum2)

print(max(sum1, sum2))

Submission status: Wrong Answer

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

def to_dict(l):
    d = {}
    for i in l:
        if i[0] in d:
            d[i[0]].append(i[1])
        else:
            d[i[0]] = [i[1]]
    return d

upper = td.copy()
upper = sorted(upper, key=lambda x: - x[1])[:k]
td = sorted(td, key=lambda x: - x[1])
dic_upper = to_dict(upper)
dic_td = to_dict(td)
kind = len(dic_td)
ans_upper = 0
c = 0
for i in dic_upper:
    ans_upper += sum(dic_upper[i])
    c += 1
ans_upper += c ** 2

ans_kind = 0
set_td = set()
c = 0
for i in td:
    if kind - len(set_td) == k - c:
        if i[0] not in set_td:
            c += 1
            ans_kind += i[1]
            set_td.add(i[0])
    else:
        c += 1
        ans_kind += i[1]
        set_td.add(i[0])
    if c == k:
        break

ans_kind += kind ** 2

print(max(ans_kind, ans_upper))

Submission status: Wrong Answer

Solution code:
N, K = list(map(int, input().split()))
data = [list(map(int, input().split())) for _ in range(N)]

data = sorted(data)
# print(data)

now = data[:K]


def count(x):
    p = len(set([a[0] for a in x]))
    point = p ** 2
    p = sum([a[1] for a in x])
    point += p
    return point

p_max = count(now)


for i in range(K, N):
    for j in range(K):
        c = now[:j] + now[j+1:] + [data[i]]
        if count(c) > p_max:
            p_max = count(c)
            now = now[:j] + now[j + 1:] + [data[i]]


print(p_max)

Submission status: Wrong Answer

Solution code:
#import collections
n, k = map(int, input().split())
c = []
for i in range(n):#h:高さ
    c.append([int(m) for m in input().split()])
c.sort(key=lambda x:(x[0], x[1]))
f = sorted(c, key=lambda x:(x[1], x[0]), reverse=True)
donyoku = 0
donset = set()
for i in range(k):
    donyoku += f[i][1]
    donset.add(f[i][0])
donyoku += len(donset)**2

t = []
d = []
for i in range(n):#h:高さ
    pi, qi = c[i][0], c[i][1]
    t.append(pi)
    d.append(qi)
spe = len(set(t))
stili = []
real = 0
reaset = set()
ws = set()
i = 0
p = [0]*n
while True:
    if f[i][0] not in ws:
        real += f[i][1]
        ws.add(f[i][0])
        p[i] = 1
        i += 1
    if i == len(donset):
        break
kosu = len(donset)
yy = real
u = 0
for i in range(len(donset) + 1, k + 1):
    while True:
        if f[i][0] not in ws:
            real += f[i][1]
            ws.add(f[i][0])
            kosu += 1
            break
        else:
            i += 1
        if i == n:
            u = 1
            break
    if u == 1:
        break
    q = 0
    while True:
        if q < n:
            if p[q] == 0:
                real += f[q][1]
                kosu += 1
                if kosu == k:
                    break
        else:
            break
        q += 1
    real += len(ws) ** 2
    stili.append(real)
    real = yy


if stili != []:
    rr = min(stili)
    print(min(donyoku,rr))
else:
    print(donyoku)

Submission status: Wrong Answer

Solution code:
import copy

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

t = [0 for _ in range(n)]
d = [0 for _ in range(n)]

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

ts = copy.copy(t)
ds = copy.copy(d)

ts.sort()
ds.sort(reverse = True)


kiso1 = 0
bonus1 = []
kiso2 = 0
bonus2 = [] #1:種類優先　2:おいしさ優先

for i in range(k):
    num = d.index(ds[i])
    kiso2 += d[num]
    bonus2.append(t[num])

point2 = kiso2 + len(set(bonus2)) ** 2

print(point2)

#print(point1)

#print(max(point1, point2))

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
N,K = [int(_) for _ in input().split()]
#dp =[[0]*(K+1) for _ in range(N+1)]

aP = sorted([[int(_) for _ in input().split()] for _ in range(N)],key=lambda x:x[1],reverse=True)
dS = set()
for x in aP[:K]:
    dS.add(x[0])
print(sum(x[1] for x in aP[:K]) + len(dS)**2)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter

n, k = map(int, input().split())
td = sorted([list(map(int, input().split())) for i in range(n)], key=itemgetter(1), reverse=True)
BP = set()
MP = 0
syurui = set()
for i in range(n):
    syurui.add(td[i][0])
syurui_n = len(syurui)
max_syurui = min(k, syurui_n)
count = 0
for i in range(n):
    t, d = td[i]
    if td[i][0] in BP and ((k != syurui_n and max_syurui == k) or k == syurui_n):
        if max_syurui ** 2 - (max_syurui - 1) ** 2 <= d:
            MP += d
            count += 1
            max_syurui -= 1
    else:
        BP.add(t)
        MP += d
        count += 1
    if count == k:
        break
print(MP + len(BP) ** 2)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
td = [list(map(int, input().split())) for _ in range(n)]
td.sort(key=lambda x:x[1],reverse=True)
a = 0
b = 0
v1 = []
v2 = []
for i in range(k):
    a += td[i][1]
    if not td[i][0] in v1:
        v1.append(td[i][0])
a += len(v1)**2
cnt = 0
for i in range(n):
    if not td[i][0] in v2:
        v2.append(td[i][0])
        b += td[i][1]
        cnt += 1
    if cnt == k:
        break
b += len(v2)**2
print(max(a,b))

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
neta = [[] for i in range(N+1)]
for i in range(N):
    t, d = map(int, input().split())
    neta[t].append(d)
for i in range(N):
    neta[i].sort(reverse = True)
neta = [i for i in neta if i != []]
neta.sort(reverse = True)
kinds = len(neta)
maxneta = 0
for i in range(len(neta)):
    maxneta = max(maxneta, len(neta[i]))
anslist = []
kinds_max = min(kinds, K)
x = K
for i in range(len(neta)):
    x -= len(neta[i])
    if x <= 0:
        kinds_min = i+1
        break
firstnetasum = [0]
for i in range(kinds):
    firstnetasum.append(firstnetasum[i] + neta[i][0])
for i in range(kinds_min, kinds_max + 1):
    ans = firstnetasum[i]
    remain = []
    for j in range(i):
        try:
            remain += neta[j][1:]
            if len(remain) > K-i:
                break
        except IndexError:
            pass
    remain.sort(reverse = True)
    ans += sum(remain[:K-i])
    ans += i**2
    anslist.append(ans)

print(max(anslist))

Submission status: Wrong Answer

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

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

a.sort(key=lambda x: x[0], reverse=True)
# おいしさ降順

ans = 0
kk = 0

bl = [False for _ in range(N + 1)]
taken = []
taken2 = [False for _ in range(N)]

for i in range(N):
    d, t = a[i]
    if not bl[t]:
        ans += d
        kk += 1
        bl[t] = True
        taken.append(i)
        taken2[i] = True
    if kk == K:
        break
ans += kk * kk

# 各種類1個ずつ最大値を取った

if kk == K:
    # 入れ替え
    taken_idx = len(taken) - 1
    while True:
        to_rm_i = taken[taken_idx]
        flg = False
        b = 0
        b -= 2 * kk - 1
        b -= a[to_rm_i][0]
        for i in range(to_rm_i):
            if taken2[i]:
                pass
            else:
                d, t = a[i]
                x = b + d
                if x > 0:
                    ans += x
                    kk -= 1
                    flg = True
                break
        if not flg: break
        taken_idx -= 1

else:
    ans -= kk * kk
    # 上から加算
    for i in range(N):
        if not taken2[i]:
            d, t = a[i]
            ans += d
            kk += 1
            if kk == K:
                break
    ans += kk * kk

print(ans)