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: Accepted

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

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

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

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

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

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

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


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

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

print(ans)

Submission status: Accepted

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

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

base = []; other = []
val = set()
yum = 0
for t, d in td[:k]:
    yum += d
    if t in val:
        other.append(d)
    val.add(t)
other.sort(reverse=True)

kouho = []
val_tmp = val.copy()
for t, d in td[k:]:
    if t not in val_tmp:
        kouho.append(d)
        val_tmp.add(t)
kouho.sort()

v = len(val)
ans = yum + (v*v)


while len(other) > 0 and len(kouho) > 0:
    yum -= other.pop()
    yum += kouho.pop()
    v += 1
    ans = max(ans, yum+(v*v))
print(ans)

Submission status: Accepted

Solution code:

import heapq

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

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

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

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
from collections import defaultdict

N, K = map(int, input().split())
TD = []

sushi = []
kinds = []


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

sushi.sort(key=lambda x: x[0], reverse=True)
sushidict = defaultdict(int)
kind_all = len(set(kinds))

cnt = 0
sushiq = []

for i in range(K):
  d, t = sushi[i]
  heappush(sushiq, (d, t))
  sushidict[t] += 1
  cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(K, N):
  d, t = sushi[i]

  if sushidict[t] > 0:
    continue

  d_min, t_min = heappop(sushiq)
  while sushidict[t_min] == 1 and sushiq:
    d_min, t_min = heappop(sushiq)

  if not sushiq:
    break

  sushidict[t_min] -= 1
  sushidict[t] += 1
  n_kind += 1
  cnt = cnt - d_min + d

  res.append(cnt+n_kind ** 2)
  if n_kind == kind_all:
    break

print(max(res))

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
import os
import itertools
import string
import heapq
import _collections
from collections import Counter
from collections import defaultdict
from functools import lru_cache
import bisect
import re
import queue


class Scanner():
    @staticmethod
    def int():
        return int(sys.stdin.readline().rstrip())

    @staticmethod
    def string():
        return sys.stdin.readline().rstrip()

    @staticmethod
    def map_int():
        return [int(x) for x in Scanner.string().split()]

    @staticmethod
    def string_list(n):
        return [input() for i in range(n)]

    @staticmethod
    def int_list_list(n):
        return [Scanner.map_int() for i in range(n)]

    @staticmethod
    def int_cols_list(n):
        return [int(input()) for i in range(n)]


class Math():
    @staticmethod
    def gcd(a, b):
        if b == 0:
            return a
        return Math.gcd(b, a % b)

    @staticmethod
    def lcm(a, b):
        return (a * b) // Math.gcd(a, b)

    @staticmethod
    def roundUp(a, b):
        return -(-a // b)

    @staticmethod
    def toUpperMultiple(a, x):
        return Math.roundUp(a, x) * x

    @staticmethod
    def toLowerMultiple(a, x):
        return (a // x) * x

    @staticmethod
    def nearPow2(n):
        if n <= 0:
            return 0
        if n & (n - 1) == 0:
            return n
        ret = 1
        while(n > 0):
            ret <<= 1
            n >>= 1
        return ret

    @staticmethod
    def sign(n):
        if n == 0:
            return 0
        if n < 0:
            return -1
        return 1

    @staticmethod
    def isPrime(n):
        if n < 2:
            return False
        if n == 2:
            return True
        if n % 2 == 0:
            return False
        d = int(n ** 0.5) + 1
        for i in range(3, d + 1, 2):
            if n % i == 0:
                return False
        return True


MOD = int(1e09) + 7
INF = int(1e15)

T_MAX = int(1e05) + 1


def main():
    # sys.stdin = open("sample.txt")
    N, K = Scanner.map_int()
    T = defaultdict(list)
    for _ in range(N):
        t, d = Scanner.map_int()
        t -= 1
        T[t].append(d)
    for v in T.values():
        v.sort(reverse=True)
    M = []
    for v in T.values():
        M.append((v[0], 1))
        for i in range(1, len(v)):
            M.append((v[i], 0))
    M.sort(reverse=True)

    q = []
    x = 0
    y = 0
    for i in range(K):
        x += M[i][0]
        y += M[i][1]
        if M[i][1] == 0:
            heapq.heappush(q, M[i][0])
    ans = x + y * y
    for i in range(K, N):
        if q == []:
            break
        if M[i][1] == 1:
            x -= heapq.heappop(q)
            x += M[i][0]
            y += 1
            ans = max(ans, x + y * y)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
import sys
from collections import *
import heapq
import math
import bisect
from itertools import permutations,accumulate,combinations,product
from fractions import gcd
def input():
    return sys.stdin.readline()[:-1]
mod=10**9+7

n,k=map(int,input().split())
td=[list(map(int,input().split())) for i in range(n)]
td.sort(key=lambda x: -x[1])

dic={}
minlst=[]
maxlst=[]
ans=0
f=0
ln=0
for i in range(n):
    t,d=td[i]
    if i<k:
        f+=d
        if not t in dic:
            dic[t]=d
        else:
            minlst.append(d)
        ln=len(dic)
    else:
        if not t in dic:
            maxlst.append(-d)
            dic[t]=d
        else:
            pass

# print(td)
# print(minlst)
# print(maxlst)
# print(f)
# print(len(dic))
# ln=len(dic)
ans=max(ans,f+pow(ln,2))
heapq.heapify(minlst)
heapq.heapify(maxlst)
while minlst and maxlst:
    f-=heapq.heappop(minlst)
    f-=heapq.heappop(maxlst)
    ln+=1
    ans=max(ans,f+pow(ln,2))
print(ans)

Submission status: Accepted

Solution code:
Num,K=map(int,input().split())
L=[[]for i in range(Num+1)]
for i in range(Num):
    t,d=map(int,input().split())
    L[t].append(d)
for i in range(1,Num+1):
    L[i].sort(reverse=True)

M=[]
m=[]
for i in range(1,Num+1):
    if len(L[i])>=1:
        M.append(L[i][0])
        if len(L[i])>=2:
            for j in range(1,len(L[i])):
                m.append(L[i][j])
M.sort(reverse=True)
m.sort(reverse=True)
MM=[0]
for i in range(len(M)):
    MM.append(MM[-1]+M[i])
mm=[0]
for i in range(len(m)):
    mm.append(mm[-1]+m[i])
#print(MM)
#print(mm)
ANS=0

def cmb(n, r, mod):
    if ( r<0 or r>n ):
        return 0
    r = min(r, n-r)
    return g1[n] * g2[r] * g2[n-r] % mod

mod = 10**9+7 #出力の制限
N = 10**5+10
g1 = [1, 1] # 元テーブル
g2 = [1, 1] #逆元テーブル
inverse = [0, 1] #逆元テーブル計算用テーブル

for i in range( 2, N + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
    inverse.append( ( -inverse[mod % i] * (mod//i) ) % mod )
    g2.append( (g2[-1] * inverse[-1]) % mod )

for i in range(1,min(K,len(M))+1):
    ANS=max(i*i+MM[i]+mm[min(len(mm)-1,K-i)],ANS)
    #print(ANS,i)
print(ANS)

Submission status: Accepted

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

Submission status: Accepted

Solution code:
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N+1)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    dt[i][1]-=1
    ans+=dt[i][0]
    if gr[dt[i][1]]==[]:
        x+=1
    gr[dt[i][1]].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    
inf=10**19
heapq.heappush(single,[inf,N])#絶対に選ばないものを入れておく
heapq.heappush(multi,[inf,N])
    
    
    
for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:#種類ボーナスも含めて小さい方を取り除く
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        tt=b[1]
        num_i[tt]+=1
        if num_i[tt]==len(gr[tt])-1:
            heapq.heappush(single,gr[tt][-1])
        else:
            heapq.heappush(multi,gr[tt][num_i[tt]])
                                       
print(ans)

Submission status: Accepted

Solution code:
#
# 　　  ⋀_⋀　 
#　　  (･ω･)  
# .／ Ｕ ∽ Ｕ＼
#  │＊　合　＊│
#  │＊　格　＊│ 
#  │＊　祈　＊│ 
#  │＊　願　＊│ 
#  │＊　　　＊│ 
#      ￣
#
import sys
sys.setrecursionlimit(10**6)
input=sys.stdin.readline
from math import floor,ceil,sqrt,factorial,hypot,log #log2ないｙｐ
from heapq import heappop, heappush, heappushpop
from collections import Counter,defaultdict,deque
from itertools import accumulate,permutations,combinations,product,combinations_with_replacement
from bisect import bisect_left,bisect_right
from copy import deepcopy
from fractions import gcd
from random import randint
inf=float('inf')
mod = 10**9+7
def pprint(*A): 
    for a in A:     print(*a,sep='\n')
def INT_(n): return int(n)-1
def MI(): return map(int,input().split())
def MF(): return map(float, input().split())
def MI_(): return map(INT_,input().split())
def LI(): return list(MI())
def LI_(): return [int(x) - 1 for x in input().split()]
def LF(): return list(MF())
def LIN(n:int): return [I() for _ in range(n)]
def LLIN(n: int): return [LI() for _ in range(n)]
def LLIN_(n: int): return [LI_() for _ in range(n)]
def LLI(): return [list(map(int, l.split() )) for l in input()]
def I(): return int(input())
def F(): return float(input())
def ST(): return input().replace('\n', '')
def main():
    N,K = MI()
    TD = LLIN(N)
    TD.sort(key = lambda x:x[1], reverse=True)

    neta = set()
    double = []
    ans = 0
    S = 0
    for i,(t,d) in enumerate(TD):
        if i < K:
            if t in neta:
                double.append(d)
            else:
                neta.add(t)
            S += d
        else:
            if t not in neta and double:
                S -= double.pop()
                S += d
                neta.add(t)
        ans = max(ans, S + len(neta)**2)
    print(ans)
if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import sys
from heapq import heappush, heappop
from collections import Counter


N, K = map(int, input().split())
#sushi = [list(map(int, x.split())) for x in sys.stdin.readlines()]
sushi = [list(map(int, input().split())) for _ in range(N)]

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

base = 0
ate = Counter()
q = []
for i in range(K):
    t, d = sushi[i]
    base += d
    heappush(q, (d, t))
    ate[t] += 1

var = len(list(ate.keys()))
candidate = [base + var ** 2]
for i in range(K, N):
    t, d = sushi[i]
    if ate[t] == 0:
        while len(q):
            dx, tx = heappop(q)
            if ate[tx] > 1:
                break
        else:
            break
        ate[t] += 1
        ate[tx] -= 1
        base += d - dx
        var += 1
        candidate.append(base + var ** 2)

print(max(candidate))

Submission status: Accepted

Solution code:
from operator import itemgetter
from itertools import accumulate


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

sushi.sort(key=itemgetter(1))
sushi.reverse()


# 各ネタ最大美味しさの寿司
a = []

# それ以外の寿司
b = []

# aにすでに入っているネタ = 1
in_a = [0]*(n+1)

for i in range(n):
    t,d = sushi[i]
    if in_a[t]==0:
        in_a[t]=1
        a.append(d)
    else:
        b.append(d)

ans = 0

a = list(accumulate([0]+a))
b = list(accumulate([0]+b))


for i in range(len(a)):
    if i<=k:
        if 0<=k-i<len(b):
            ans = max(a[i]+b[k-i]+i**2,ans)
print(ans)

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
l = [tuple(map(int, input().split())) for _ in range(n)]
l.sort(key=lambda x:-x[1])
dic = dict()
tr = list()
curd = 0
curt = 0
for t,d in l[:k]:
  curd += d
  if t not in dic:
    curt += 1
    dic[t] = d
  elif dic[t] < d:
    tr.append(dic[t])
    dic[t] = d
  else:
    tr.append(d)
ans = curt**2 + curd
tr.sort(reverse=True)
for t,d in l[k:]:
  if t in dic: continue
  if not tr: break
  curt += 1
  dic[t] = d
  curd += d - tr.pop()
  cur = curt**2 + curd
  if ans < cur: ans = cur
print(ans)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())  # N個の寿司からK個
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)  # 美味しさで降順ソート
kind = set()
stack = []
ans = 0
for t, d in sushi[:K]:
    if t in kind:
        stack.append(d)
    else:
        kind.add(t)
    ans += d

ans += len(kind)**2
point = ans
for t, d in sushi[K:]:
    if len(stack) == 0:  # K個の寿司全てが1つずつしか存在しないなら交換しようがない
        break

    # 交換を行う
    if t in kind:
        continue
    else:
        point -= stack.pop()
        point += d
        point += 2*len(kind) + 1  # 平方の数列の階差
        kind.add(t)
    ans = max(ans, point)
print(ans)

Submission status: Accepted

Solution code:
def examA():
    L = LI(); L.sort()
    ans = L[0]*L[1]//2
    print(ans)
    return

def examB():
    def fun(s):
        if s%2==0:
            return s//2
        else:
            return s*3+1
    S = I()
    d = defaultdict(bool)
    d[S] = True; cur = 1
    while(True):
        cur +=1
        S = fun(S)
        if d[S]:
            break
        d[S] = True
    print(cur)
    return

def examC():
    N = I()
    H = LI()
    cur = 0
    ans =0
    for i in range(N):
        if cur<H[i]:
            ans += H[i]-cur
        cur = H[i]
    print(ans)
    return

def examD():
    N, K = LI()
    TS = [LI() for _ in range(N)]
    choice = []
    keep = []
    d = defaultdict(int)
    TS.sort(key=lambda x:x[1],reverse=True)
    for i,(t,s) in enumerate(TS):
        if d[t]==0:
            if i < K:
                d[t] = -1
                choice.append(s)
            else:
                d[t] = s
        else:
            if i<K:
                keep.append(s)
    keep.sort()
    D = sorted(d.values(),reverse=True)
#    print(choice); print(keep); print(D)
    cur = len(choice)
    now = cur**2 +sum(choice)+sum(keep)
    ans = copy.deepcopy(now)
    for i,k in enumerate(keep):
        if D[i]<0:
            break
        now += ((cur+i)*2+1) + (D[i]-k)
        ans = max(ans,now)
#        print(now)
    print(ans)
    return

import sys,copy,bisect,itertools,heapq,math
from heapq import heappop,heappush,heapify
from collections import Counter,defaultdict,deque
def I(): return int(sys.stdin.readline())
def LI(): return list(map(int,sys.stdin.readline().split()))
def LSI(): return list(map(str,sys.stdin.readline().split()))
def LS(): return sys.stdin.readline().split()
def SI(): return sys.stdin.readline().strip()
global mod,mod2,inf,alphabet,ALPHABET, _ep
mod = 10**9 + 7
mod2 = 998244353
inf = 10**18
_ep = 1e-9
alphabet = [chr(ord('a') + i) for i in range(26)]
ALPHABET = [chr(ord('A') + i) for i in range(26)]

if __name__ == '__main__':
    examD()

Submission status: Accepted

Solution code:
n,k=[int(x) for x in input().split()]
A=[]
for i in range(n):
  a,b=[int(x) for x in input().split()]
  A.append([b,a,True])

A.sort(reverse=True)

typ=set()

for i in range(n):
  if A[i][1] in typ:
    A[i][2]=False
  else:
    typ.add(A[i][1])

num=0

for j in range(k):
  if A[j][2]==True:
    #print(j)
    num+=1

ans=0
for i in range(k):
  ans+=A[i][0]
ans+=num**2



curr=ans
nxt = k
lst = k-1

while(True):

	while(nxt < n and A[nxt][2] == False):
		nxt += 1
	while(lst >= 0 and A[lst][2] == True):
		lst -= 1
	if not(nxt < n and lst >= 0):
		break

	curr += A[nxt][0] - A[lst][0]
	curr +=  (num+1) ** 2 - num**2


	num+=1
	nxt+= 1
	lst  -= 1

	ans=max(ans,curr)




print(ans)

Submission status: Accepted

Solution code:
def p_d():
    N, K = map(int, input().split())
    inf = 10 ** 18
    sushi = [[-inf] for _ in range(N)]
    max_kind = -1
    for _ in range(N):
        t, d = map(int, input().split())
        sushi[t - 1].append(d)
        if max_kind < t:
            max_kind = t
    for i in range(max_kind):
        sushi[i].sort(reverse=True)
    sushi.sort(key=lambda x: x[0], reverse=True)

    que = []
    res = 0
    for i in range(K):
        res += sushi[i][0]
        for j in sushi[i][1:]:
            que.append(j)
    que.sort()
    cur = res
    res += K ** 2
    for var in range(1, K)[::-1]:
        a = que.pop()
        b = sushi[var][0]
        cur += a - b
        res = max(res, cur + var ** 2)
    print(res)


if __name__ == '__main__':
    p_d()

Submission status: Accepted

Solution code:
import sys
from heapq import heappush, heappop

n, k, *td = map(int, sys.stdin.read().split())
td = list(zip(*[iter(td)] * 2))

def main():
    td.sort(key=lambda x: x[1], reverse=True)
    delis = [[] for _ in range(n+1)]
    cnt = [0] * (n + 1)
    base = 0
    x = 0
    toppings = set()
    for t, d in td[:k]:
        toppings.add(t)
        if cnt[t] == 0:
            x += 1
        cnt[t] += 1
        heappush(delis[t], d)
        base += d
    score = base + x * x

    cand_to_remove = []
    for t in range(1, n+1):
        if cnt[t] >= 2:
            heappush(cand_to_remove, (heappop(delis[t]), t))
            cnt[t] -= 1

    for t, d in td[k:]:
        if not cand_to_remove:
            break
        if t in toppings:
            continue
        r_d, r_t = heappop(cand_to_remove)
        toppings.add(t)
        base -= r_d
        cnt[t] += 1
        heappush(delis[t], d)
        base += d
        x += 1
        if cnt[r_t] >= 2:
            heappush(cand_to_remove, (heappop(delis[r_t]), r_t))
            cnt[r_t] -= 1
        score = max(score, base + x * x)

    return score
        
if __name__ == '__main__':
    ans = main()
    print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
ds = [[] for _ in range(n)] # ネタの種類t(<=n)別においしさdのリスト
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    ds[t].append(d)

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    # 各ネタでおいしさのsort
    y1.append(ds[i][-1])    # 各ネタで最大のおいしさをpopしy1に追加
    ds[i].pop()
    y0.extend(ds[i])        # 各ネタで2番目以降のおいしさのリストをy0に追加
y0.sort(reverse=True)   # ネタの種類が増えないおいしさのリスト（各ネタで2番目以降）を降順sort
y1.sort(reverse=True)   # ネタの種類が増えるおいしさ（各ネタで最大）のリストを降順sort

ans = 0
# Y = max(0, k-len(y0))   # Y: ネタの種類の最小値（kまで増やしながら貪欲するため）
Y = min(k, len(y1))   # Y: ネタの種類の最大値（減らしながら貪欲するため）
X = 0
for i in range(Y):
    X += y1[i]      # Y個 ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に加算
for i in range(k-Y):
    X += y0[i]      # k-Y個 ネタの種類が増えないおいしさを大きい順に加算

while True:
    ans = max(ans, X+Y*Y)
    # if (Y>=k or Y>=len(y1)):   # Yをkまで増やしながら貪欲
    #     break
    # X += y1[Y]      # ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に1つ追加
    # X -= y0[k-Y-1]  # ネタの種類が増えない大きさを小さい順に1つ除外
    # Y += 1
    if (Y<=0 or k-Y>=len(y0)):   # Yを減らしながら貪欲
        break
    X -= y1[Y-1]      # ネタの種類が増えるおいしさ（各ネタで最大）を小さい順に1つ除外
    X += y0[k-Y]  # ネタの種類が増えない大きさを大きい順に1つ追加
    Y -= 1
print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
ds = [[] for _ in range(n)] # ネタの種類t(<=n)別においしさdのリスト
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    ds[t].append(d)

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    # 各ネタでおいしさのsort
    y1.append(ds[i][-1])    # 各ネタで最大のおいしさをpopしy1に追加
    ds[i].pop()
    y0.extend(ds[i])        # 各ネタで2番目以降のおいしさのリストをy0に追加
y0.sort(reverse=True)   # ネタの種類が増えないおいしさのリスト（各ネタで2番目以降）を降順sort
y1.sort(reverse=True)   # ネタの種類が増えるおいしさ（各ネタで最大）のリストを降順sort

ans = 0
Y = max(0, k-len(y0))   # Y: ネタの種類の最小値（kまで増やしながら貪欲するため）
X = 0
for i in range(Y):
    X += y1[i]      # Y個 ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に加算
for i in range(k-Y):
    X += y0[i]      # k-Y個 ネタの種類が増えないおいしさを大きい順に加算

while True:
    ans = max(ans, X+Y*Y)
    if (Y>=k or Y>=len(y1)):   # Yをkまで増やしながら貪欲
        break
    X += y1[Y]      # ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に1つ追加
    X -= y0[k-Y-1]  # ネタの種類が増えない大きさを小さい順に1つ除外
    Y += 1
print(ans)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
li=[list(map(int,input().split())) for i in range(n)]
gr=[0]*(n+1) #各グループの出現回数をメモする配列

from operator import itemgetter
li=sorted(li, key=itemgetter(1), reverse=True)

A=[] #大きい方からK個以内の配列で、2個以上あるやつの2番目以降のやつの大きさを記録

for i in range(k):
  if gr[li[i][0]]>=1:
    A.append(li[i][1])
  gr[li[i][0]]+=1

A.reverse()

B=[] #既に同じグループが登場してしまってない子の大きさを記録
for i in range(k,n):
  if gr[li[i][0]]==0:
    B.append(li[i][1])
  gr[li[i][0]]+=1

#現時点での種類数とポイントを測定
sp=set()
ans=0
for i in range(k):
  ans+=li[i][1]
  sp.add(li[i][0])
sp=len(sp)
ans+=sp*sp

koukan=[] #その交換で、累積でどの程度増減するか
zouka=0
for a,b in zip(A,B):
  zouka+=(sp+1)**2-sp**2-(a-b)
  #print(a,b,sp,zouka)
  koukan.append(zouka)
  sp+=1

#koukanの中で最大の増加を探して、ansに+する。
# print(A,B)
# print(koukan)
if koukan:
  ans+=max(0,max(koukan))

print(ans)

Submission status: Accepted

Solution code:
import sys
from heapq import heappush,heappop,heapify
sys.setrecursionlimit(10**9)
INF=10**18
MOD=10**9+7
def input():
    return sys.stdin.readline().rstrip()

def main():
    N,K=map(int,input().split())
    sushies=[]
    for i in range(N):
        t,d=map(int,input().split())
        sushies.append((d,t-1))
    sushies.sort()
    b=[False]*N
    overlappedsushi=[]
    ans=0
    n=0
    for i in range(K):
        sushi=sushies.pop()
        if b[sushi[1]]:
            heappush(overlappedsushi,sushi)
            ans+=sushi[0]
        else:
            b[sushi[1]]=True
            ans+=sushi[0]
            n+=1
    ans+=n**2
    tmp=ans
    while overlappedsushi:
        worstsushi=heappop(overlappedsushi)
        tmp-=worstsushi[0]
        while sushies:
            bestsushi=sushies.pop()
            if not b[bestsushi[1]]:
                break
        else:
            break
        n+=1
        tmp+=bestsushi[0]
        tmp+=n*2-1
        b[bestsushi[1]]=True
        ans=max(ans,tmp)
    print(ans)

    

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda p: -p[1])
neta = set()
oisi = netan = 0
amari = []
for t, d in TD[:K]:
    oisi += d
    if t not in neta:
        neta.add(t)
        netan += 1
    else:
        amari.append(d)
ans = oisi + netan * netan
for t, d in TD[K:]:
    if netan == K or not amari:
        break
    if t not in neta:
        neta.add(t)
        netan += 1
        oisi += d - amari.pop()
        ans = max(ans, oisi + netan * netan)
print(ans)

Submission status: Accepted

Solution code:
import os
import sys
from collections import deque
from operator import itemgetter

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353

N, K = list(map(int, sys.stdin.buffer.readline().split()))
TD = [list(map(int, sys.stdin.buffer.readline().split())) for _ in range(N)]

t_seen = [False] * (N + 1)
sushi_seen = [False] * N
TD.sort(key=itemgetter(1), reverse=True)

delicious = 0
que = deque()
for i, (t, d) in enumerate(TD):
    if len(que) >= K:
        break
    if not t_seen[t]:
        que.append(i)
        delicious += d
        t_seen[t] = True
        sushi_seen[i] = True
for i, (t, d) in enumerate(TD):
    if len(que) >= K:
        break
    if not sushi_seen[i]:
        que.appendleft(i)
        delicious += d
        sushi_seen[i] = True

idx = i
type_cnt = sum(t_seen)
ans = delicious + type_cnt ** 2
while idx < que[-1]:
    while idx < N and sushi_seen[idx]:
        idx += 1
    if idx >= N:
        break
    que.appendleft(idx)
    delicious -= TD[que.pop()][1]
    delicious += TD[idx][1]
    type_cnt -= 1
    score = delicious + type_cnt ** 2
    ans = max(ans, score)
    sushi_seen[idx] = True
print(ans)

Submission status: Accepted

Solution code:
import heapq
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)
neta = [0] * (n+1)
hq = []
kind = 0
ans = 0
for t,d in td[:k]:
  if neta[t]:
    hq.append(d)
  else:
    kind += 1
  neta[t] += 1
  ans += d
ans += kind ** 2
heapq.heapify(hq)
now = k
new_ans = ans
while hq and now < n:
  t,d = td[now][0],td[now][1]
  if not neta[t]:
    MIN = heapq.heappop(hq)
    new_ans += d - MIN + 2 * kind + 1
    ans = max(ans,new_ans)
    kind += 1
    neta[t] += 1
  now += 1	
print(ans)

Submission status: Accepted

Solution code:
import itertools

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

td = {}

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

top = []
rest = []

for k in td.keys():
    td[k].sort()
    top.append(td[k][-1])
    rest += td[k][:-1]

top.sort()
rest.sort()

cumtop = list(itertools.accumulate(reversed(top)))
cumrest = list(itertools.accumulate(reversed(rest)))

ret = 0

for i in range(min(len(top), K)):
    tmp = (i+1)*(i+1)
    tmp += cumtop[i]
    if i == K-1:
        pass
    else:
        if (K-i-2) < 0 or (K-i-2) > len(cumrest)-1:
            continue
        tmp += cumrest[(K-i-2)]
    ret = max(tmp, ret)

print(ret)

Submission status: Accepted

Solution code:
import heapq

N,K = map(int, input().split())
h = []
max_t = 0
for _ in range(N):
    t, d = map(int, input().split())
    heapq.heappush(h, (-d, t))
    max_t = max(max_t, t)

L = [0]*max_t
num = 0
h2 = []
for _ in range(K):
    temp = heapq.heappop(h)
    L[temp[1]-1] += 1
    num -= temp[0]
    heapq.heappush(h2, (-temp[0], temp[1]))
    
cnt = 0
for i in range(max_t):
    if L[i] > 0:
        cnt += 1

ans = num + cnt**2
while len(h) > 0 and len(h2) > 0:
    temp = heapq.heappop(h)
    if L[temp[1]-1] > 0:
        continue
    while L[h2[0][1]-1] <= 1:
        heapq.heappop(h2)
        if not h2:
            break
    if not h2:
        break
    temp2 = heapq.heappop(h2)    
    num -= temp2[0] + temp[0]
    L[temp2[1]-1] -= 1
    L[temp[1]-1] += 1
    cnt += 1
    ans = max(ans, num+cnt**2)

print(ans)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
inf = 10**18
sushi = [[-inf] for _ in range(N)]
for i in range(N):
    t, d = map(int, input().split())
    sushi[t-1].append(d)
for i in range(N):
    sushi[i].sort(reverse = True)
sushi.sort(key = lambda x:x[0], reverse = True)
que = []
res = 0
for i in range(K):
    res+=sushi[i][0]
    for j in sushi[i][1:]:
        que.append(j)
que.sort()
cur = res
res+=K**2
for var in range(1, K)[::-1]:
    a = que.pop()
    b = sushi[var][0]
    cur+=a-b
    res = max(res, cur+var**2)
print(res)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
N, M = map(int, input().split())

s = [tuple(map(int, input().split())) for _ in range(N)]
s.sort(key = lambda x: -x[1])
table = [(-1, -1)] * (N + 1)
c = 0
for i in range(N):
  if table[s[i][0]] == (-1, -1):
    table[s[i][0]] = (i, s[i][1])
    c += 1

table.sort(key = lambda x: -x[1])
res = [0] * (min(c, M) + 1)
vis = set()
p = M
for i in range(M): res[0] += s[i][1]
for x in range(1, min(c, M) + 1):
  res[x] = res[x - 1] + pow(x, 2) - pow(x - 1, 2)
  if table[x - 1][0] >= M:
    res[x] += table[x - 1][1]
    p -= 1
    while p in vis:
      p -= 1
    res[x] -= s[p][1]
  vis.add(table[x - 1][0])
  #print(p)
print(max(res))