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:
N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for i in range(N)]
point = [0 for i in range(K+1)]

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

eaten_sushi = [sushi[i] for i in range(K)]
eaten_variation = [s[0] for s in eaten_sushi]
variation_count = dict([(i, eaten_variation.count(i)) for i in range(1, N+1)])
variation_number = len(set(eaten_variation))
base = sum([s[1] for s in eaten_sushi])
bonus = variation_number ** 2
point[variation_number] = base + bonus


for i in range(K, N):
    if variation_count[sushi[i][0]] == 1:
        continue

    for j in range(K-1, 0, -1):
        if variation_count[eaten_sushi[j][0]] > 0:
            continue
        base = base - sushi[j][1] + sushi[i][1]
        variation_count[sushi[j][0]] += -1
        variation_count[sushi[i][0]] = 1
        eaten_sushi.pop(j)
        eaten_sushi.append(sushi[i])
        variation_number += 1
        bonus = variation_number ** 2
        point[variation_number] = base + bonus
        break

print(max(point))

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(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
if len(Y) >= k:
    x = 0
    y = sum(Y)
    Y = Y[:k]
    ans = sum(Y) + k ** 2
else:
    x = sum(X[:len(Y) - k])
    y = sum(Y)
    ans = sum(X[:len(Y) - k]) + sum(Y) + len(Y) ** 2
    X = X[len(Y) - k:]
X = X[::-1]
t = ans
len_y = len(Y)
while X and Y:
    len_y -= 1
    x += X.pop()
    y -= Y.pop()
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

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(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
if len(Y) >= k:
    x = 0
    y = sum(Y)
    Y = Y[:k]
    ans = sum(Y) + k ** 2
else:
    x = sum(X[:len(Y) - k])
    y = sum(Y)
    ans = sum(X[:len(Y) - k]) + sum(Y) + k ** 2
    X = X[len(Y) - k:]
X = X[::-1]
t = ans
len_y = len(Y)
while X and Y:
    len_y -= 1
    x += X.pop()
    y -= Y.pop()
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

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(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
if len(Y) >= k:
    x = 0
    y = sum(Y)
    Y = Y[:k]
    ans = sum(Y) + k ** 2
else:
    x = sum(X[len(Y) - k:])
    y = sum(Y)
    ans = sum(X[:len(Y) - k]) + sum(Y) + k ** 2
    X = X[len(Y) - k:]
X = X[::-1]
t = ans
len_y = len(Y)
while X and Y:
    len_y -= 1
    x += X.pop()
    y -= Y.pop()
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
n,k=map(int,input().split())
z=[0,1]
for i in range(1,n):
    z.append(z[i]+2)
s=sorted([list(map(int,input().split()))for _ in range(n)],key=lambda x:-x[1])
q=[]
v=0
ans=0
l=[0]*-~n
for a in s[:k]:
    if not l[a[0]]:
        v+=1
        ans+=z[v]
    else:
        heapq.heappush(q,a[1])
    l[a[0]]=1
    ans+=a[1]
tem=ans
for a in s[k:]:
    if not q:break
    if l[a[0]]:continue
    l[a[0]]=1
    v+=1
    t=heapq.heappop(q)
    tem=tem-t+a[1]+z[v]
    ans=max(ans,tem)
print(tem)

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
group = [[] for i in range(N)]
for i in range(N):
    t,d = map(int,input().split())
    t -= 1
    group[t].append(d)
for e in group:
    e.sort()
    e.reverse()
head = []
remain = []
for e in group:
    if e == []:
        continue
    head.append(e[0])
    for i in range(1,len(e)):
        remain.append(e[i])
head.sort()
head.reverse()
remain.sort()
remain.reverse()
ans = -float("inf")
for i in range(1,len(head)+1):
    if K - i > len(remain):
        continue
    deli = sum(head[:i]) + sum(remain[:K-i])
    kind = i**2
    res = deli + kind
    if res > ans:
        ans = res
print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())

top = [-1]*n # t毎に最高のdを格納するlist
sub = [] # それ以外のdを格納するlist
cst = [0] # cumsum top
css = [0] # cumsum sub

for _ in range(n):
    t, d = (int(i) for i in input().split())
    t -= 1
    if top[t] == -1:
        top[t] = d
    elif d > top[t]:
        top[t] = d
        sub.append(top[t])
    else:
        sub.append(d)

top.sort(reverse=True)
sub.sort(reverse=True)

for i in range(len(top)): # カテゴリ内最大のdが高い順に累積
    if top[i] != -1:
        cst.append(cst[-1] + top[i])

for i in range(len(sub)): # それ以外の全てのカテゴリの寿司のdが高い順に累積
    css.append(css[-1] + sub[i])

po = 0
for i in range(1, k+1):
    if len(cst) >= i+1 and len(css) >= k-i+1:
        po = max(cst[i] + css[k-i] + i*i, po)

print(po)

Submission status: Wrong Answer

Solution code:
import copy
import itertools

def funcF(p):
    return p

def funcG(p):
    return funcF(p) + p*p


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

# print(tdList)
tdList.sort(key=lambda x: x[1])
tdList.reverse()

S = [[0 for i in range(K)] for j in range(K+1)]
f = [0] * (K+1)
g = [0] * (K+1)
# tmpS=[i for i in range(K)]
x=[]
tmpF=0
for i in range(K):
    tmpF += tdList[i][1]
    if tdList[i][0] not in x:
        x.append(tdList[i][0])
a=len(x)
f[a] = tmpF
g[a] = f[a] + a*a
S[a] = [i for i in range(K)]

tmpT = []
for i in range(K):
    tmpT.append(tdList[S[a][i]][0])
tmpK = K-1
tmpS = K-1
for i in range(a+1, K+1):
    S[i] = copy.deepcopy(S[i-1])
    f[i] = f[i-1]
    flag = False

    for j in range(tmpK, -1, -1):
        if tmpT.count(tmpT[j]) == 1:
            pass
        else:
            f[i] -= tdList[S[i][j]][1]
            del(tmpT[j])
            del(S[i][j])
            tmpK = j-1
            flag = True
            break
    if flag == False:
        break

    flag = False
    for j in range(tmpS+1, N):
        if tdList[j][0] not in tmpT:
            flag = True
            tmpS = j
            S[i].append(j)
            f[i] += tdList[j][1]
            g[i] = f[i] + i*i
            break
    if flag == False:
        break

# print(tdList)
# print(S)
# print(f)
# print(g)
# print(g.index(max(g)))
print(max(g))



# combList = list(itertools.combinations(range(N), K))
# manzokudo=[0]*len(combList)
# for combID in range(len(combList)):
#     comb = combList[combID]
#     x=[]
#     for i in range(K):
#         manzokudo[combID] += d[comb[i]]
#         if t[comb[i]] in x:
#             pass
#         else:
#             x.append(t[comb[i]])
#     manzokudo[combID] += len(x)*len(x)
#
# print(max(manzokudo))

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
td = [list(map(int, input().split())) for i in range(N)]
td = sorted(td, key=lambda x:x[1], reverse=True)
td_dict = {}
for key in td:
    td_dict[key[0]] = []
for key in td:
    td_dict[key[0]].append(key[1])
max_list = []
else_list = []
for key in td_dict.keys():
    tmp = sorted(td_dict[key], reverse=True)
    max_list.append(tmp[0])
    if len(tmp) > 1:
        else_list += tmp[1:]    
max_list = sorted(max_list, reverse=True)
else_list = sorted(else_list, reverse=True)
point = 0
for k in range(1, K+1):
    tmp = k**2 + sum(max_list[:k]) + sum(else_list[:K-k])
    if tmp > point:
        point = tmp
print(point)

Submission status: Wrong Answer

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

td = [list(map(int, input().split())) for i in range(N)]
td = sorted(td, key=lambda x:x[1], reverse=True)
td_dict = {}
for key in td:
    td_dict[key[0]] = []
for key in td:
    td_dict[key[0]].append(key[1])
for key in td:
    td_dict[key[0]] = sorted(td_dict[key[0]], reverse=True)
max_list = []
else_list = []
for key in td_dict.keys():
    max_list.append(td_dict[key][0])
    if len(td_dict[key]) > 1:
        else_list += td_dict[key][1:]
max_list = sorted(max_list, reverse=True)
else_list = sorted(else_list, reverse=True)
point = 0
for k in range(1, K+1):
    tmp = k**2 + sum(max_list[:k]) + sum(else_list[:K-k])
    if tmp > point:
        point = tmp
print(point)

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
group = [[] for i in range(N)]
for i in range(N):
    t,d = map(int,input().split())
    t -= 1
    group[t].append(d)
for e in group:
    e.sort()
    e.reverse()
head = []
remain = []
for e in group:
    if e == []:
        continue
    head.append(e[0])
    for i in range(1,len(e)):
        remain.append(e[i])
head.sort()
head.reverse()
remain.sort()
remain.reverse()
ans = -float("inf")
for i in range(K-len(remain),len(head)+1):
    deli = sum(head[:i]) + sum(remain[:K-i])
    kind = i**2
    res = deli + kind
    if res > ans:
        ans = res
print(ans)

Submission status: Wrong Answer

Solution code:

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

def get_duplicate_list(seq,index):
    seen = []
    return [x[index] for x in seq if not seen.append(x[index]) and seen.count(x[index]) == 2]

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

sorted_sushi=sorted(sushi,key= lambda x:x[1])
#print(sorted_sushi)
taberu=[]
for i in range(K):
    taberu.append(sorted_sushi.pop())

sorted_taberu = sorted ( taberu, key=lambda x: -x[ 1 ] )
maxscore=0
for i in range(N):
    # 重複要素をだす
    if len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu]) >maxscore:
        maxscore=len(set([x[0] for x in sorted_taberu])) ** 2 + sum ( [x[1] for x in sorted_taberu])
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sorted_taberu = sorted ( sorted_taberu, key=lambda x: -x[ 1 ],reverse=True )
    dp=get_duplicate_list(sorted_taberu,0)
    '''
    print('pop taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    if not dp:
        if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
            maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                [ x[ 1 ] for x in sorted_taberu ] )
        continue
    else:
        sb=False
        for i in range(len(sorted_taberu)):
            for s in dp:
                if sorted_taberu[i][0]==s:
                    sorted_sushi.append(sorted_taberu.pop(i))
                    sb=True
                    break
            if sb:
                break
    sorted_sushi = sorted ( sorted_sushi, key=lambda x: x[ 1 ] )
    sb=False
    '''
    print('push taberu \n')
    print('sushi')
    print(sorted_sushi)
    print('taberu')
    print(sorted_taberu)
    print('\n\n')
    '''
    for i in range(len(sorted_sushi)):
        setted_sushi=set([x[0] for x in sorted_sushi])
        setted_taberu=set([x[0] for x in sorted_taberu])

        seki=setted_sushi & setted_taberu

        if not seki:
            sorted_taberu.append(sorted_sushi.pop())
            if len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] ) > maxscore:
                maxscore = len ( set ( [ x[ 0 ] for x in sorted_taberu ] ) ) ** 2 + sum (
                    [ x[ 1 ] for x in sorted_taberu ] )
                sb=True
                break

        for i in range(len(sorted_sushi)):
            for s in list(seki):
                if sorted_sushi[i][0]!=seki:
                    sb=True
                    sorted_taberu.append ( sorted_sushi.pop (i ) )
                    break
            if sb:
                break
        if sb:
            break




print (maxscore)
exit ( 0 )

Submission status: Wrong Answer

Solution code:
import sys

sys.setrecursionlimit(200000)


def input():
    return sys.stdin.readline()[:-1]


def ii(t: type = int):
    return t(input())


def il(t: type = int):
    return list(map(t, input().split()))


def imi(N: int, t: type = int):
    return [ii(t) for _ in range(N)]


def iml(N: int, t: type = int):
    return [il(t) for _ in range(N)]


def solve():
    N, K = il()
    TD = iml(N)
    TD.sort(key=lambda x: x[1], reverse=True)
    TD = list(zip(*TD))
    kind = list(TD[0][:K])[::-1]
    neta = list(TD[1][:K])[::-1]
    ans = sum(neta) + len(set(kind)) ** 2
    for i in range(K, N):
        t, d = TD[0][i], TD[1][i]
        if t in kind:
            continue
        kind.append(t)
        neta.append(d)
        ans = max(ans, sum(neta[-K:]) + len(set(kind[-K:])) ** 2)
    return ans


if __name__ == "__main__":
    print(solve())

Submission status: Wrong Answer

Solution code:
import numpy as np
from itertools import chain


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

TDP = []

for n in range(N):
    TDP.append(list(map(int, input().split())) + [0])

TDP = np.array(TDP)
TDP = TDP[np.lexsort((TDP[:, 1] * -1, TDP[:, 0]))]

# u, c = np.unique(TDP[:, 0], return_counts=True) after numpy 1.9
t = np.bincount(TDP[:, 0])
u = np.nonzero(t)[0]
c = t[u]

TDP[:, 2] = list(chain.from_iterable([list(range(c[i])) for i in range(len(c))]))

TDP = TDP[np.lexsort((TDP[:, 1] * -1, TDP[:, 2]))]

L = TDP[:len(u), 0].tolist()

R = []

def isin(element, test_elements, assume_unique=False, invert=False):
    element = np.asarray(element)
    return np.in1d(element, test_elements, assume_unique=assume_unique, invert=invert).reshape(element.shape)

for p in range(1, K+1):
    R.append(sum(TDP[isin(TDP[:, 0], L[:p])][:K, 1]) + p ** 2)

print(max(R))

Submission status: Wrong Answer

Solution code:
N,K = list(map(int, input().split()))
s = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x:-x[1])
 
#単純に
res = sorted(s[:K] ,key=lambda x: x[0], reverse=True)
p = sum([i[1] for i in res]) + len(set([i[0] for i in res]))**2

#入れ替えて
o1 = s[K:]
d = set([i[0] for i in o1]) - set([i[0] for i in res])
o2 = []

for i in range(N-K):
    if o1[i][0] in d:
        if i > 0:
            if o1[i][0] != o1[i-1][0]:
                o2.append(o1[i])
        else:
            o2.append(o1[i])

r1 = []
r2 = []
for i in range(K):
    if i > 0:
        if res[i][0] != res[i-1][0]:
            r1.append(res[i])
        else:
            r2.append(res[i])
    else:
        r1.append(res[i])
        
r2 = sorted(r2 ,key=lambda x: x[1], reverse=True)
o2 = sorted(o2 ,key=lambda x: x[1], reverse=True)

for i in range(1,len(o2)+1):
    x = r1 + r2[:-i] + o2[:i]
    if p < sum([i[1] for i in x]) + len(set([i[0] for i in x]))**2:
        p = sum([i[1] for i in x]) + len(set([i[0] for i in x]))**2

print(p)

Submission status: Wrong Answer

Solution code:
from collections import Counter
from heapq import heappush, heappop
N, K = map(int, input().split())
sushi_list = []
for i in range(N):
    t, d = map(int, input().split())
    sushi_list.append((t, d))


sushi_list.sort(key=lambda x: x[1], reverse=True)
g = []
S_x = sushi_list[:K]
c = Counter()
f_p = 0

for t, d in S_x:
    c[t] += 1
    f_p += d

Q = []
R = []

for sushi in sushi_list:
    if sushi[0] in c.keys():
        if c[sushi[0]] > 1:
            heappush(Q, sushi[1])
    else:
        heappush(R, -sushi[1])
p = len(c)
ans = f_p + p**2

while True:
    if len(Q) == 0 or len(R) == 0:
        break
    min_q = heappop(Q)
    max_r = -heappop(R)
    p += 1
    f_p = f_p + max_r - min_q
    ans = max(ans, f_p + p**2)

print(ans)

Submission status: Wrong Answer

Solution code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][1])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            if ans > num:
                break
            else:
                ans = num
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][1])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            ans = max(ans, num)
            break
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
def main():
    N, K = map(int, input().split())
    td = []
    for _ in range(N):
        ti, di = map(int, input().split())
        td.append([ti, di])
    td.sort(key=lambda x: -x[1])
    num = 0
    kind = set()
    changable = []
    for i in range(K):
        num += td[i][1]
        if td[i][0] in kind:
            changable.append(td[i][0])
        else:
            kind.add(td[i][0])
    num += len(kind) ** 2
    ans = num

    for i in range(K, N):
        if not changable:
            break
        if td[i][0] not in kind:
            num += td[i][1] + 2 * len(kind) + 1 - changable.pop()
            kind.add(td[i][0])
            ans = max(ans, num)


if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
ans = 0
aru = []
nai = []
spe = set()
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans+=anstmp[1]
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = ans + sum(aru) + s**2

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans+=nai[i][1]
        s += 1
    m = max(m, ans + sum(aru) + s**2)
print(m)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = set()
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = sum(ans) + sum(aru) + s**2
#print(ans, aru, s)

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans.append(nai[i][1])
        s += 1
        #print(ans, aru, s)
    m = max(m, sum(ans) + sum(aru) + s**2)
print(m)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
ans = []
aru = []
nai = []
spe = set()
f = []
m = 0
for i in range(n):
    t, d = map(int, input().split())
    nai.append([t,d])
nai = sorted(nai, key=lambda x: x[1], reverse = True)
for i in range(k):
    anstmp = nai.pop(0)
    if(anstmp[0] not in spe):
        ans.append(anstmp[1])
        spe.add(anstmp[0])
    else:
        aru.append(anstmp[1])

s = len(spe)
m = sum(ans) + sum(aru) + s**2

for i in range(len(nai)):
    if(len(aru) == 0): break
    if(nai[i][0] not in spe):
        aru.pop()
        ans.append(nai[i][1])
        s += 1
    m = max(m, sum(ans) + sum(aru) + (s + i)**2)
print(m)

Submission status: Wrong Answer

Solution code:


def solve(N, K, tds):
    tds.sort(reverse=True, key=lambda x: x[1])
    tmp_max = sum(t[1] for t in tds[:K])
    various = len(set([t[0] for t in tds[:K]]))

    has = tds[:K]
    v = {t[0]: 0 for t in has}
    for t in has:
        v[t[0]] += 1
    has.sort(key=lambda x: x[1])

    nothas = tds[K:]
    nothas.sort(key=lambda x: x[1])

    mm = tmp_max + various ** 2
    for _ in range(0, 10000):
        target = None
        for t in has:
            if v[t[0]] > 1:
                v[t[0]] -= 1
                tmp_max -= t[1]
                target = t
                break
        if target:
            has.remove(target)
        else:
            break

        target = None
        for t in nothas:
            if t[0] not in v:
                v[t[0]] = 1
                tmp_max += t[1]
                target = t
                break
        if target:
            nothas.remove(target)
        else:
            break

        various += 1
        mm = max(mm, tmp_max + various ** 2)

    return mm

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    tds = [tuple(map(int, input().split(" "))) for _ in range(N)]
    print(solve(N, K, tds))

Submission status: Wrong Answer

Solution code:

def solve(N, K, tds):
    tds.sort(reverse=True, key=lambda x: x[1])
    tmp_max = sum(t[1] for t in tds[:K])
    various = len(set([t[0] for t in tds[:K]]))

    has = tds[:K]
    v = {t[0]: 0 for t in has}
    for t in has:
        v[t[0]] += 1
    has.sort(key=lambda x: x[1])

    nothas = tds[K:]
    nothas.sort(key=lambda x: x[1])

    mm = tmp_max + various ** 2
    for _ in range(0, 10000):
        if all(v[t[0]] == 1 for t in has):
            break

        if all(t[0] in v for t in nothas):
            break

        for t in has:
            if v[t[0]] > 1:
                has.remove(t)
                v[t[0]] -= 1
                tmp_max -= t[1]
                break
        for t in nothas:
            if t[0] not in v:
                has.append(t)
                v[t[0]] = 1
                tmp_max += t[1]
                break
        various += 1
        mm = max(mm, tmp_max + various ** 2)

    return mm

if __name__ == "__main__":
    N, K = tuple(map(int, input().split(" ")))
    tds = [tuple(map(int, input().split(" "))) for _ in range(N)]
    print(solve(N, K, tds))

Submission status: Wrong Answer

Solution code:
import sys
from collections import Counter
from itertools import combinations

import numpy as np

sys.setrecursionlimit(200000)


def input():
    return sys.stdin.readline()[:-1]


def ii(t: type = int):
    return t(input())


def il(t: type = int):
    return list(map(t, input().split()))


def imi(N: int, t: type = int):
    return [ii(t) for _ in range(N)]


def iml(N: int, t: type = int):
    return [il(t) for _ in range(N)]


def solve():
    N, K = il()
    TD = np.asarray(iml(N))
    TD = np.sort(TD, axis=1)
    ans = 0
    for i in range(N - K):
        ans = max(ans, np.sum(TD[i : i + K][:, 1]) + len(set(TD[i : i + K][:, 0])) ** 2)
    return ans


if __name__ == "__main__":
    print(solve())

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)

l = min(K,len(res2))
p2 = 0
for i in range(l+1):
    res4 = res2[:i] + res3[:K-i] 
    if p2 < sum([i[1] for i in res4]) + len(set([i[0] for i in res4]))**2:
        p2 = sum([i[1] for i in res4]) + len(set([i[0] for i in res4]))**2

print(max(p1,p2))

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]
            already_eat.add(sushi[i][1])
        print max_score + len(already_eat) * len(already_eat)
        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=map(int,input().split())

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

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

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
delind=K-1

ans = res+x**2

while delind>=0 and repl<=N-1:
    while delind>=0:
        t,d=sushi[delind]
        if ate[t]>1:
            ate[t]-=1
            res-=d
            delind-=1
            break
        delind-=1

    while repl<=N-1:
        t,d=sushi[repl]
        if ate[t]==0:
            ate[t]+=1
            x+=1
            res+=d
            repl+=1
            break
        repl+=1
    ans=max(ans,res+x**2)

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)

    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 + len(already_eat) * len(already_eat)
        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()