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:
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])
                    kind+=1
                    break
            if flag:
                break
    print(ans+max(ans_list))

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)
ma_syurui = max_syurui
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 ma_syurui ** 2 - (ma_syurui - 1) ** 2 <= d:
            MP += d
            count += 1
            ma_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:
import sys
fin = sys.stdin.readline
from heapq import heappop, heappush
from collections import defaultdict


N, K = [int(elem) for elem in fin().split()]
sushi_info = sorted([tuple(int(elem) for elem in fin().split()) for _ in range(N)], key=lambda x: -x[1])


cur_eaten = []
cur_rest = []
count = defaultdict(int)
for t, d in sushi_info[:K]:
    heappush(cur_eaten, (d, t))
    count[t] += 1


for t, d in sushi_info[K:]:
    heappush(cur_rest, (-d, t))

num_distinctive_kinds = sum(v >= 1 for v in count.values())
satisfactions = [sum(d for _, d in sushi_info[:K]) + num_distinctive_kinds**2]

cur_num_kinds = num_distinctive_kinds
for cur_num_kinds in range(num_distinctive_kinds + 1, N + 1):
    while cur_eaten and count[cur_eaten[0][1]] <= 1:
        heappop(cur_eaten)
    while cur_rest and count[cur_rest[0][1]] > 0:
        heappop(cur_rest)
    if (not cur_eaten) or (not cur_rest):
        break
    ditched_d, ditched_t = heappop(cur_eaten)
    added_d, added_t = heappop(cur_rest)
    added_d *= -1
    count[ditched_t] -= 1
    count[added_d] += 1
    # cur_num_kinds += 1
    satisfactions.append(satisfactions[-1] - ditched_d + added_d + cur_num_kinds * 2 - 1)

# print(satisfactions)
print(max(satisfactions))
# num_kinds = max(t for t, _ in sushi_info)

# is_eaten = [False] * num_kinds
# satisfaction = 0

# for _ in range(K):

Submission status: Wrong Answer

Solution code:
N,K=map(int,input().split())
L=[list(map(int,input().split())) for i in range(N)]
L.sort(key=lambda L:(L[1],L[0]), reverse=True)
delicious=0
kind=[]
for i in range(K):
    delicious+=L[i][1]
    kind.append(L[i][0])
num=len(set(kind))
point=delicious+num**2
print(point)

Submission status: Wrong Answer

Solution code:
import sys

N,K=map(int, sys.stdin.readline().split()) 
T=[ map(int, sys.stdin.readline().split()) for i in range(N) ]

T=sorted(T, key=lambda x:(x[0],-x[1]))

prev=0
for i  in range(N):
	x=T[i][0]
	if prev!=x:
		T[i][0]=prev+1
	else:
		T[i][0]=prev+1
	prev=x		
		


A=[ (0,0) ]
dsum=0

for t,d in T:
	dsum+=d
	A.append( (t,dsum) )

	
ans=float("-inf")
for i in range(N-K+1):
	ans=max(ans, (A[i+K][0]-A[i][0])**2 + (A[i+K][1]-A[i][1]) )

print ans

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)
    if len(s)==len(d):
        return tmp
    up=n[1]-1
    down=n[1]
    while len(s)!=len(d) and 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:
N, K = map(int, input().split())
X = []
for i in range(N):
    x, y = map(int, input().split())
    X.append([i, x, y])

X = sorted(X, key = lambda x:x[2])[::-1]

p = 0
L = [0] * (N+1)
bonus = 0
for i in range(K):
    p += X[i][2]
    if L[X[i][1]] == 0:
        bonus += 1
        
    L[X[i][1]] += 1

p += bonus ** 2
ma = p
i = K-1
j = K
while True:
    while j < N and L[X[j][1]]:
        j += 1
    if j == N:
        break
    else:
        while i >= 0 and L[X[i][1]] <= 1:
            i -= 1
        if i < 0:
            break
        else:
            L[X[i][1]] -= 1
            L[X[j][1]] += 1
            p -= X[i][2]
            p += X[j][2]
            p += bonus * 2 + 1
            bonus += 1
            ma = max(ma, p)
            
print(ma)

Submission status: Wrong Answer

Solution code:
#coding:utf-8
import copy
N, K = map(int,input().split())
T = []
D = []
ls_td = []
for n in range(N):
    t, d = map(int,input().split())
    ls_td.append([t,d])
ls_td = sorted(ls_td, key=lambda x: x[1])[::-1]
head = ls_td[:K]
tail = ls_td[K:]
def get_set(mat):
    st = set()
    for t,d in mat:
        st.add(t)
    return st
def score(mat):
    r = 0
    ls_t = set()
    for t,d in mat:
        ls_t.add(t)
        r += d
    r += len(ls_t)**2
    return r

def search_del(mat):
    idx = -1
    st = set()
    for i,(t,d) in enumerate(mat):
        if t in st:
            idx = i
        st.add(t)
    return idx
st = get_set(head)
current_score = score(head)
idx = search_del(head)
for t,d in tail:
    if idx == -1:
        break

    if t in st:
        pass
    else:
        new = copy.deepcopy(head)
        new[idx] = t,d
        candidate_score = score(new)
        if current_score < candidate_score:
            head = new
            st.add(t)
            current_score = candidate_score
            idx = search_del(head)
print (current_score)

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(len(li)+1)]

for i in range(min(len(li),len(l))):
    d[i+1][0] += b[i] - a[i]
    d[i+1][1] = min(k,t+i+1)
ans = 0
for i in range(len(li)+1):
    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:
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
                    bl[t] = True
                    taken2[i] = True
                    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)

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)
    if len(s)==len(d):
        return tmp
    up=n[1]-1
    down=n[1]
    while len(s)!=len(d) and 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
                
        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
    return tmp
print(func())

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)
    if len(s)==len(d):
        return tmp
    up=n[1]-1
    down=n[1]
    while len(s)!=len(d) and 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
                
        while down<len(k):
            if k[down][0] not in d:
                d[k[down][0]]=0
                result+=k[down][1]
                tmp=max(tmp,result+len(d)*len(d))
            down +=1
    return tmp
print(func())

Submission status: Wrong Answer

Solution code:

N,K = map(int,input().split())
td = []
td1 = [[] for _ in range(N+1)]
for i in range(N):
    t,d = map(int,input().split())
    td.append([t,d,0])
td = sorted(td,key = lambda x:x[1])[::-1]
ans = 0
s = set()
for i in range(K):
    ans += td[i][1]
    s.add(td[i][0])
ans += len(s)**2
ans1 = 0
n = set()
length = 0
count = 0
for _ in range(N):
    n = set()
    for i in range(N):
        t,d,f = td[i]
        if f == 0:
            if t not in n and count < K:
                td[i][2] = 1
                n.add(t)
                ans1 += d
                count += 1
                continue
    length = max(length,len(n))
    if count >= K :
        break
ans1 += length**2
print(max(ans,ans1))

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])
                    kind+=1
            if flag:
                break
    print(ans+max(ans_list))

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
    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())

M = [[None for k in range(K+1)] for n in range(N+1)]
T = [[None for k in range(K+1)] for n in range(N+1)]
D = [[None for k in range(K+1)] for n in range(N+1)]

for n in range(N+1):
    if n == 0:
        M[n] = [0] * (K+1)
        T[n] = [[] for k in range(K+2)]
        D[n] = [0] * (K+1)
        continue
    t, d  = map(int, input().split())
    for k in range(K+1):
        if k == 0:
            M[n][k] = 0
            T[n][k] = []
            D[n][k] = 0
        elif k > n:
            M[n][k] = M[n][k-1]
            T[n][k] = T[n][k-1]
            D[n][k] = D[n][k-1]
        else:
            T_add = list(set(T[n-1][k-1] + [t]))
            D_add = D[n-1][k-1] + d
            M_add = D_add + len(set(T_add)) ** 2
            if M_add > M[n-1][k]:
                M[n][k] = M_add
                T[n][k] = T_add
                D[n][k] = D_add
            elif M_add == M[n-1][k]:
                if len(T_add) < len(T[n-1][k]):
                    M[n][k] = M_add
                    T[n][k] = T_add
                    D[n][k] = D_add
                else:
                    M[n][k] = M[n-1][k]
                    T[n][k] = T[n-1][k]
                    D[n][k] = D[n-1][k]
            else:
                M[n][k] = M[n-1][k]
                T[n][k] = T[n-1][k]
                D[n][k] = D[n-1][k]

print(M[N][K])

Submission status: Wrong Answer

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  A.append(list(map(int,input().split())))

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
l = 0
le = N - 1
k = False
D = []
while l < K:
  while True:
    j += 1
    if j >= le:
      j = 0
      k = True
    if k or not A[j][0] in D:
      ans2 += A[j][1]
      D.append(A[j][0])
      A[j],A[le] = A[le],A[j]
      l += 1
      le -= 1
      break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

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)
    if len(s)==len(d):
        return tmp
    up=n[1]-1
    down=n[1]
    while len(s)!=len(d) and 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
                
        while down<len(k):
            if k[down][0] not in d:
                d[k[down][0]]=0
                result+=k[down][1]
                tmp=max(tmp,result+len(d)*len(d))
            down +=1
    return tmp
print(func())

Submission status: Wrong Answer

Solution code:
#coding:utf-8
import copy
N, K = map(int,input().split())
T = []
D = []
ls_td = []
for n in range(N):
    t, d = map(int,input().split())
    T.append(t)
    D.append(d)
    ls_td.append([t,d])
#print (T,D)
ls_td = sorted(ls_td, key=lambda x: x[1])[::-1]
head = ls_td[:K]
tail = ls_td[K:]
#print (head)
#print (tail)
def get_set(mat):
    st = set()
    for t,d in mat:
        st.add(t)
    return st
def score(mat):
    r = 0
    ls_t = []
    for t,d in mat:
        ls_t.append(t)
        r += d
    r += len(set(ls_t))**2
    return r

def search_del(mat):
    idx = -1
    d_del = -1
    st = set()
    for i,(t,d) in enumerate(mat):
        if t in st:
            idx = i
            d_del = d
        st.add(t)
    return idx
st = get_set(head)
current_score = score(head)
idx = search_del(head)
for t,d in tail:
    if idx == -1:
        break

    if t in st:
        pass
    else:
        new = copy.deepcopy(head)
        #print (new[idx])
        new[idx][0] = t
        new[idx][1] = d
        candidate_score = score(new)
        if current_score < candidate_score:
            #head.pop(idx)
            #head.append([t,d])
            head = new
            st.add(t)
            current_score = candidate_score
            idx = search_del(head)
print (current_score)

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)
    if len(s)==len(d):
        return tmp
    up=n[1]-1
    down=n[1]
    while len(s)!=len(d):

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

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(len(li)+1)]

for i in range(min(len(li),len(l))):
    d[i+1][0] += b[i] - a[i]
    d[i+1][1] = t+i+1
ans = 0
for i in range(len(li)+1):
    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:
N, K = map(int, input().split())

M = [[None for k in range(K+1)] for n in range(N+1)]
T = [[None for k in range(K+1)] for n in range(N+1)]
D = [[None for k in range(K+1)] for n in range(N+1)]

for n in range(N+1):
    if n == 0:
        M[n] = [0] * (K+1)
        T[n] = [[] for k in range(K+2)]
        D[n] = [0] * (K+1)
        continue
    t, d  = map(int, input().split())
    for k in range(K+1):
        if k == 0:
            M[n][k] = 0
            T[n][k] = []
            D[n][k] = 0
        elif k > n:
            M[n][k] = M[n-1][k-1]
            T[n][k] = T[n-1][k-1]
            D[n][k] = D[n-1][k-1]
        else:
            T_add = list(set(T[n-1][k-1] + [t]))
            D_add = D[n-1][k-1] +  d
            M_add = D_add + len(set(T_add)) ** 2
            if M_add > M[n-1][k]:
                M[n][k] = M_add
                T[n][k] = T_add
                D[n][k] = D_add
            elif M_add == M[n-1][k]:
                if len(T_add) < len(T[n-1][k]):
                    M[n][k] = M_add
                    T[n][k] = T_add
                    D[n][k] = D_add
                else:
                    M[n][k] = M[n-1][k]
                    T[n][k] = T[n-1][k]
                    D[n][k] = D[n-1][k]
            else:
                M[n][k] = M[n-1][k]
                T[n][k] = T[n-1][k]
                D[n][k] = D[n-1][k]

print(M[N][K])

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 = []
    remaincount = K-i
    for j in range(i):
        try:
            a = neta[j][1:]
            remaincount -= len(a)
            remain[len(remain):] = a
            if remaincount <= 0:
                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:
if __name__ == '__main__':
    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:
                if c[td[check][0]]==1:
                    check -=1
                    if check == 0:
                        flag =True
                        break
                else:
                    c[td[check][0]]-=1
                    ans_list.append(2*kind+1+d-td[check][1])
                    kind+=1
            if flag:
                break
    print(ans+max(ans_list))

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))
count = 0
sum1=0
for i in range(n):
    b = np.asarray(list(map(int, input().split())))
    a = np.r_[a,b.reshape(1,-1)]
    t = a[:,0]
    d = a[:,1]
#print(a,t,d) # [a1, a2, a3, ..., aN]
mt = [0,]
for _ in range(k):
    maxIndex = np.argmax(d)#[i for i, x in enumerate(d) if x == max(d)]
    sum1+=d[maxIndex]
    d[maxIndex] = 0
    #print(maxIndex)
    flag = 0
    for mtmp in mt:
        if mtmp == t[maxIndex]:
            flag = 1
            break
    
    if flag == 0:
        mt.append(t[maxIndex])
        #print(mt)
        count += 1

print(int(sum1+count*count))

Submission status: Wrong Answer

Solution code:
import sys

N,K=map(int, sys.stdin.readline().split()) 
T=[ map(int, sys.stdin.readline().split()) for i in range(N) ]

T=sorted(T, key=lambda x:(x[0],-x[1]))


A=[ (0,0) ]
dsum=0

for t,d in T:
	dsum+=d
	A.append( (t,dsum) )

	
ans=float("-inf")
for i in range(N-K+1):
	ans=max(ans, (A[i+K][0]-A[i][0])**2 + (A[i+K][1]-A[i][1]) )

print ans

Submission status: Wrong Answer

Solution code:
import sys
fin = sys.stdin.readline
from heapq import heappop, heappush
from collections import defaultdict


N, K = [int(elem) for elem in fin().split()]
sushi_info = sorted(tuple(tuple(int(elem) for elem in fin().split()) for _ in range(N)), key=lambda x: -x[1])

cur_eaten = []
cur_rest = []
count = defaultdict(int)
for t, d in sushi_info[:K]:
    heappush(cur_eaten, (d, t))
    count[t] += 1
for t, d in sushi_info[K:]:
    heappush(cur_rest, (-d, t))

num_distinctive_kinds = sum(v >= 1 for v in count.values())
satisfactions = [sum(d for _, d in sushi_info[:K]) + num_distinctive_kinds**2]

cur_num_kinds = num_distinctive_kinds
for cur_num_kinds in range(num_distinctive_kinds + 1, K + 1):
    while cur_eaten and count[cur_eaten[0][1]] <= 1:
        heappop(cur_eaten)
    while cur_rest and count[cur_rest[0][1]] > 0:
        heappop(cur_rest)
    if (not cur_eaten) or (not cur_rest):
        break
    ditched_d, ditched_t = heappop(cur_eaten)
    added_d, added_t = heappop(cur_rest)
    added_d *= -1
    count[ditched_t] -= 1
    count[added_d] += 1
    # cur_num_kinds += 1
    satisfactions.append(satisfactions[-1] - ditched_d + added_d + cur_num_kinds * 2 - 1)

# print(satisfactions)
print(max(satisfactions))