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:
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(100000)


def read_int():
    return int(input())


def read_int_n():
    return list(map(int, input().split()))


def read_float():
    return float(input())


def read_float_n():
    return list(map(float, input().split()))


def read_str():
    return input().strip()


def read_str_n():
    return list(map(str, input().split()))


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap


@mt
def slv(N, K, TD):
    TD.sort(key=lambda x: x[1], reverse=True)

    c = Counter()
    base = 0
    rc = []
    for i in range(K):
        t, k = TD[i]
        c[t] += 1
        base += k
        if c[t] > 1:
            rc.append((t, k))

    mi = i
    ans = [base + len(c)**2]
    while True:
        if len(rc) == 0:
            break
        rt, rk = rc.pop()
        cc = len(c)
        d = (cc+1)**2 - cc**2
        for i in range(mi+1, len(TD)):
            if TD[i][0] not in c:
                c[TD[i][0]] += 1
                c[rt] -= 1
                base -= rk
                base += TD[i][1]
                mi = i
                ans.append(base + len(c)**2)
                break

        else:
            break
    return max(ans)


def main():
    # N = read_int()
    N, K = read_int_n()
    TD = [read_int_n() for _ in range(N)]
    print(slv(N, K, TD))


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
def various_sushi(N: int, K: int, TD: list) -> int:
    _TD = sorted(TD, key=lambda x: x[1], reverse=True)

    tops, others = [], []
    tops_sum, others_sum = 0, 0
    kind = set()
    res = 0
    q = []
    for t, d in _TD[:K]:
        if t not in kind:
            kind.add(t)
        else:
            q.append(d)
        res += d

    max_d = res + len(kind) ** 2
    temp = max_d
    count = len(q)
    for t, d in _TD[K:]:
        if count == 0:
            break
        if t not in kind:
            temp += -q.pop() + d + len(kind) * 2 + 1
            kind.add(t)
            max_d = max(max_d, temp)
            count -= 1

    return max_d


if __name__ == "__main__":
    N, K = [int(s) for s in input().split()]
    TD = [tuple(int(s) for s in input().split()) for _ in range(N)]

    ans = various_sushi(N, K, TD)
    print(ans)

Submission status: Accepted

Solution code:
from collections import deque
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]
res = 0
td.sort(key = lambda x:-x[1])
kind = set()
queue = deque()
for t, d in td[:k]:
    if t in kind:
        queue.append(d)
    else:
        kind.add(t)
    res += d
ans = res + len(kind)**2
point = ans
count = len(queue)
for t, d in td[k:n]:
    if count == 0:
        break
    if not t in kind:
        point += -queue.pop() + d + len(kind) * 2 + 1
        kind.add(t)
        ans = max(ans, point)
        count -= 1
print(ans)

Submission status: Accepted

Solution code:
N,K = (int(x) for x in input().split())
t_arr = [0]*(N+1)
d_arr = [0]*(N+1)
for i in range(N):
    t,d = tuple(int(x) for x in input().split())
    t_arr[i+1] = t
    d_arr[i+1] = d

td_index = list(range(1,N+1))
td_index.sort(key=lambda x:d_arr[x], reverse=True)

n_dm_set = set()
small_arr = []
d_dm = 0
for i in td_index[:K]:
    t = t_arr[i]
    d = d_arr[i]
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

current_max = len(n_dm_set)**2+d_dm

n_x_set = set(n_dm_set)
d_x = d_dm
nt_x = len(n_dm_set)
for i in td_index[K:]:
    if not small_arr:
        break
    t = t_arr[i]
    d = d_arr[i]
    if t in n_x_set:
        continue
    n_x_set.add(t)
    d_x += d
    d_x -= small_arr.pop()
    nt_x += 1
    current_max = max(current_max, nt_x**2+d_x)

print(current_max)

Submission status: Accepted

Solution code:
import sys
input=sys.stdin.readline
sys.setrecursionlimit(10**9)
n,k=map(int,input().split())
Sushi=[list(map(int,input().split())) for i in range(n)]
Sushi=sorted(Sushi,key=lambda x: x[1],reverse=True)
ans=0
A=[]
kind=set()
for t,d in Sushi[:k]:
    ans+=d
    if t in kind:
        A.append(d)
    else:
        kind.add(t)
ans+=len(kind)**2
tmp=ans
for t,d in Sushi[k:]:
    if not A:
        break
    if not t in kind:
        tmp+=d-A.pop()+2*len(kind)+1
        kind.add(t)
        ans=max(ans,tmp)
print(ans)

Submission status: Accepted

Solution code:
from operator import itemgetter

n, k = map(int, input().split())
td = [tuple(int(x) for x in input().split()) for _ in range(n)]

a = sorted(td, key=itemgetter(1), reverse=True)

ans = [0]
s = set()
t = []
for i in range(k):
    ans[0] += a[i][1]
    if a[i][0] not in s:
        s.add(a[i][0])
    else:
        t.append(a[i][1])

p = len(s)
q = len(t) - 1

if q < 0:
    print(ans[0] + k * k)
    exit()

v = []
for i in range(k, n):
    if a[i][0] not in s:
        s.add(a[i][0])
        v.append(a[i][1])

ans[0] += p * p
d = 0

for x in v:
    ans.append(ans[d] + x - t[q] + (p + 1) * (p + 1) - p * p)
    p += 1
    d += 1
    q -= 1
    if q < 0:
        break

print(max(ans))

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
point=0
stack=[]
kind=set()
for t,d in TD[:K]:
    point+=d
    if t in kind:stack.append(d)
    else:kind.add(t)
point=point+len(kind)**2
ans=point
for t,d in TD[K:]:
    if not stack:break
    if not t in kind:
        point+=d-stack.pop()+2*len(kind)+1
        kind.add(t)
        ans=max(ans,point)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop

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

dsum = 0
used = {}
q = []
for _ in range(K):
    d, t = heappop(h)
    if t not in used:
        used[t] = 1
    else:
        heappush(q, (-d, t))
    dsum += -d
res = len(used)**2 + dsum

while h and q:
    d, t = heappop(h)
    if t not in used:
        used[t] = 1
        m, _ = heappop(q)
        dsum = dsum - m - d
        res = max(res, len(used)**2 + dsum)
print(res)

Submission status: Accepted

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

td = sorted(td, key = lambda x: x[1], reverse = True)
eat_types = set()
sushi_dup = []
value = 0
for t, d in td[:k]:
    if not t in eat_types:
        eat_types.add(t)
    else:
        sushi_dup.append(d)
    value += d

ans = value + len(eat_types)**2
for t, d in td[k:]:
    if len(sushi_dup) == 0:
        break
    
    if t in eat_types:
        continue
    else:
        eat_types.add(t)
        r = sushi_dup.pop()
        value += d - r
        ans = max(ans, value + len(eat_types)**2)
print(ans)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
lis = []
val = 0
ans = 0
s = set()
dub = []

for _ in range(N):
    lis.append(tuple(map(int,input().split())))

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

for t,v in lis[:K]:

    if t in s:
        dub.append(v)
    s.add(t)
    val += v

ans = val + len(s) ** 2

for t,v in lis[K:]:
    if len(dub) == 0:
        break
    if t in s:
        continue

    s.add(t)
    r = dub.pop()
    val += v - r
    ans = max(ans,val+len(s) ** 2)

print(ans)

Submission status: Accepted

Solution code:
#参考 https://atcoder.jp/contests/abc116/submissions/4059000
n, k = map(int,input().split())
# 価値の大きい順にソート, おいしさが前に来るようにした
sushi = sorted(list(map(int, input().split()))[::-1] for i in range(n))[::-1]
s = set()
# 二個以上あるものの集合があれば良い,かつ最小順
q = []
kiso = 0
for i in range(k):
    d, t = sushi[i]
    if t in s:
        q.append(d)
    else:
        s.add(t)
    kiso += d
ans = kiso + len(s)**2

# 減らせる ∧ 増やせる
i = k
while i < n and q:
    d, t = sushi[i]
    if t not in s:
        s.add(t)
        kiso += -q.pop()+d
        ans = max(ans, kiso+len(s)**2)
        continue
    i+=1
print(ans)

Submission status: Accepted

Solution code:
import sys
import heapq
input = sys.stdin.readline
def main():
    n,k = map(int, input().split())
    sushi = [tuple(map(int, input().split())) for _ in range(n)]
    sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
    species = set()
    val = 0
    dup = []
    # dの大きい順に計算する
    for t,d in sushi[:k]:
        # ネタが被った寿司の点を持っておく(この部分を他のネタの寿司と入れ替える)
        if t in species:
            dup.append(d)
        species.add(t)
        val += d
    ans = val + len(species)**2

    for t,d in sushi[k:]:
        if len(dup)==0:
            break
        # 同じネタはすでに最大になるように選んでいるのでパス
        # 種類は減らない
        if t in species:
            continue
        species.add(t)
        removal = dup.pop()
        val += d - removal
        ans = max(ans, val + len(species)**2)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
import heapq

N,K=map(int,input().split())
TD=[[int(x) for x in input().split()] for y in range(N)]

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

ans=0
neta=set([])
rep_neta=[]

for i in range(K):
    a,b=TD[i]
    if not (a in neta):
        neta.add(a)
        ans+=b
    else:
        ans+=b
        rep_neta.append(b)

ans+=len(neta)**2

cnt=ans

for i in range(K,N):
    if len(rep_neta) == 0 :
        break
    else:
        a,b=TD[i]
        if not (a in neta):
            cnt+=2*len(neta)+1
            cnt+=b-rep_neta.pop(-1)
            neta.add(a)
            ans=max(ans,cnt)

print(ans)

Submission status: Accepted

Solution code:
N, K = [int(x) for x in input().split()]
td = sorted([[int(x) for x in input().split()] for _ in range(N)],
            key=lambda x: x[1], reverse=True)
INF = 10**9
# print(td)

first = []
second = []
kind = set()
for i in range(K):
    if td[i][0] not in kind:
        first.append(td[i])
        kind.add(td[i][0])
    else:
        second.append(td[i])

yummy = sum([x[1] for x in first + second])
best_ans = yummy + len(kind)**2

for k in range(K, N):
    if len(second) == 0:
        break
    if td[k][0] in kind:
        continue

    a = second.pop(-1)
    yummy -= a[1]
    kind.add(td[k][0])
    first.append(td[k])
    yummy += td[k][1]
    ans = yummy + len(kind)**2
    # print(first, second, ans)
    if best_ans < ans:
        best_ans = ans
print(best_ans)

Submission status: Accepted

Solution code:
#i!/usr/bin/env python
# coding: utf-8

import sys
input = sys.stdin.readline

def main():
    N, K = list(map(int, input().split()))
    l = []
    for _ in range(N):
        t, d = list(map(int, input().split()))
        l.append((d, t))
    l.sort(reverse=True)
    first = []
    second = []
    t_set = set()
    for k in range(K):
        if l[k][1] not in t_set:
            first.append(l[k])
            t_set.add(l[k][1])
        else:
            second.append(l[k])
    yummy = sum(e[0] for e in first) + sum(e[0] for e in second)
    best_ans = yummy + len(t_set)**2
    for k in range(K, N):
        if len(second) == 0:
            break
        if l[k][1] in t_set:
            continue
        sushi = second.pop(-1)
        yummy -= sushi[0]
        yummy += l[k][0]
        t_set.add(l[k][1])
        first.append(l[k])
        ans = yummy + len(t_set)**2
        if ans  > best_ans:
            best_ans = ans
    print(best_ans)


if __name__ == '__main__':
    main()

Submission status: Accepted

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

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

D = []
T = {}
d = 0

for i in range(K):
  D.append(td[i])
  d += td[i][1]
  t = td[i][0]
  if t in T:
    T[t] += 1
  else:
    T[t] = 1

D.sort(key = lambda x: x[1])

cnt = 0
flag = False
Ans = d + pow(len(T), 2)
for i in range(K, N):
  if td[i][0] in T:
    continue
  while T[D[cnt][0]] == 1:
    cnt += 1
    if cnt == K:
      flag = True
      break
  if flag:
    break

  T[td[i][0]] = 1
  T[D[cnt][0]] -= 1
  d -= D[cnt][1] - td[i][1]
  
  t = d + pow(len(T), 2)
  if Ans < t:
    Ans = t
  cnt += 1
  if cnt == K:
    break

print(Ans)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
TD = [[int(x) for x in input().split()] for _ in range(N)]
TD.sort(key = lambda x: x[1], reverse=True)

T = 0
neta = set()
doublecounts = []
ans = []

for i in range(K):
    T += TD[i][1]    
    if not(TD[i][0] in neta):    
        neta.add(TD[i][0])       
    else:      
        doublecounts.append(TD[i][1])    
            
ans.append(T+len(neta)**2)

for i in range(K,N):
    if(len(doublecounts)==0):
        break
    if not(TD[i][0] in neta):
        T -= doublecounts.pop(-1)
        T += TD[i][1]
        neta.add(TD[i][0])   
        ans.append(T+len(neta)**2)
print(max(ans))

Submission status: Accepted

Solution code:
from collections import deque


def main():
    N, K = map(int, input().split())
    sushis = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        sushis[i] = (t, d)
    sushis.sort(key=lambda a: (-a[1], a[0]))
    tmp = 0
    types = set()
    stack = deque()
    ans = 0
    for i in range(len(sushis)):
        t, d = sushis[i]
        if i < K:
            if t in types:
                stack.append(d)
            types.add(t)
            tmp += d
        else:
            if t not in types and len(stack):
                types.add(t)
                tmp += d
                tmp -= stack.pop()
        ans = max(ans, tmp + len(types)**2)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, vari


selected = sushis[:K]
unSelected = sushis[K:]

p, vari = calcPoint(selected)
cand = [p]

va = len(vari)
v = [0]*(N+1)
for i, _ in selected:
    v[i] += 1

can_del = []
for i, d in reversed(selected):
    if v[i] >= 2:
        v[i] -= 1
        can_del.append(d)

can_in = []
for i, d in unSelected:
    if not i in vari:
        vari.add(i)
        can_in.append(d)

#print("del->", can_del, can_in)
for od, nd in zip(can_del, can_in):
    nextp = p-od+nd-(va)**2+(va+1)**2
    p = nextp
    cand.append(p)
    va += 1
print(max(cand))

Submission status: Accepted

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



def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))
#-----------------------------------------

li.sort(reverse=True,key=lambda x:x[1])
s2=min(find_num_neta(list(range(n)),li),k)
s1=find_num_neta(list(range(k)),li)
oishisa=sum([li[i][1] for i in range(k)])
score=oishisa+s1**2

#-----------------------------------------
li.sort(key=lambda x:x[0])
best_list=[]
not_best_list=[]


m=-1
for i in li:
    if m==i[0]:
        not_best_list.append(i)
    else:
        best_list.append(i)
        m=i[0]
        
not_best_list.sort(reverse=True,key=lambda x:x[1])
best_list.sort(reverse=True,key=lambda x:x[1])



 
 
for i in range(s1+1,s2+1):
    oishisa=oishisa+best_list[i-1][1]-not_best_list[k-i][1]
    
    new_score=i**2+oishisa
    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
S=[] #[ネタ種類、おいしさ]のリスト
for i in range(n):
    S.append(list(map(int,input().split())))
S.sort(reverse=1,key=lambda x:x[1])

netano_min=min(S[j][0] for j in range(n))
for i in range(n):
    S[i][0]-=netano_min-1

neta=max(S[j][0] for j in range(n)) #ネタの種類数

F=[0 for i in range(neta+1)] #ネタi種のときの美味しさの和の最大値のリスト

ma=sum(S[i][1] for i in range(k))

neta_s=len(set([S[i][0] for i in range(k)]))

netano_min=min(S[j][0] for j in range(n))

F[neta_s]=ma

#print(S)
#print(F)
#print(ma)

netalist=[0 for i in range(n+1)] #i種類目のネタを使った個数
keseru=[] #ネタに重複がある寿司のリスト(消してもいい寿司)
for i in range(k):
    if netalist[S[i][0]]>=1 :
        keseru.append(S[i])
        
    netalist[S[i][0]]+=1

#print(keseru)
#print(netalist)

right=len(keseru)-1
left=k
i=neta_s+1
while len(keseru)>0 and i<=k and right>=0 and left<=n-1 :
    while left<=n-1 and netalist[S[left][0]]!=0 :
        left+=1
    if left==n :
        break
    netalist[S[left][0]]+=1
    F[i]=F[i-1]+S[left][1]-keseru[right][1]
    right-=1
    left+=1
    i+=1

#print(F)
for j in range(neta_s,i+1-1) :
    F[j]+=j**2

print(max(F))

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
netas_oishisas = []
for _ in range(n):
    netas_oishisas.append(tuple(map(int, input().split())))
netas_oishisas.sort(key = lambda no: no[1], reverse = True)
tabetaneta = set()
second = []
greedyscore = 0
for i in range(k):
    neta, oishisa = netas_oishisas[i]
    if neta in tabetaneta:
        second.append(oishisa)
    else:
        tabetaneta.add(neta)
    greedyscore += oishisa
greedyscore += len(tabetaneta) ** 2
maxscore = greedyscore
currentscore = greedyscore
second.reverse()
secondpointer = 0
secondn = len(second)
for i in range(k, n):
    if secondpointer >= secondn:
        break
    neta, oishisa = netas_oishisas[i]
    if neta not in tabetaneta:
        currentscore = currentscore - second[secondpointer] + 2*len(tabetaneta) + 1 + oishisa
        maxscore = max(maxscore, currentscore)
        secondpointer += 1
        tabetaneta.add(neta)
print(maxscore)

Submission status: Accepted

Solution code:
import sys

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    S = TD[:K]
    result = base_point + type_bonus**2
    
    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        if type_count_dict[s[0]] >= 2:
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            type_bonus += 1
            ridx += 1
            result = max(type_bonus**2 + base_point, result)
    print(result)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
#  import math, string, itertools, fractions, heapq, collections, re,  array, bisect, sys, random, time, copy, functools
import sys
#  import numpy as np
#  from queue import PriorityQueue

sys.setrecursionlimit(10**7)
inf = 10 ** 20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ddn = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]


def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    N, K = LI()
    T = []
    D = []
    TD = []
    [TD.append(LI()) for i in range(N)]
    # sort desc
    #  TD = np.array(TD)
    #  TD = TD[np.argsort(TD[:, 1])[::-1]]
    TD = sorted(TD, key=lambda l:l[1], reverse=True)
    #  print(TD)

    #  print(TD)
    # 種類のカウントをもつ
    type_count_dict = {}
    base_point = 0
    type_bonus = 0
    remaining_items = []
    for idx, i in enumerate(TD):
        if not type_count_dict.get(i[0]):
            type_count_dict[i[0]] = 0
        if idx < K:
            if type_count_dict[i[0]] == 0:
                type_bonus += 1
            type_count_dict[i[0]] += 1
            base_point += i[1]
        else:
            if type_count_dict[i[0]] == 0:
                remaining_items.append(i[1])
                type_count_dict[i[0]] += 1

    #  print(type_count_dict)

    S = TD[:K]
    #  print('S', S)
    #  for i in S:
    #      type_count_dict[i[0]] += 1
    #      base_point += i[1]

    #  print('type_count', type_count_dict)
    #  base_point = sum(i[1] for i in S)
    #  type_bonus = len(set([i[0] for i in S]))
    result = base_point + type_bonus**2
    #  results = [result]
    #  print('base_point, type_bonus, result', base_point, type_bonus**2, result)

    #  remaining_items = []
    #  remaining_items = PriorityQueue()
    #  done_types = []
    #  print('type_count_dict', type_count_dict)
    #  for i in TD[K:]:
    #      if i[0] in done_types:
    #          continue
    #      if type_count_dict[i[0]] == 0:
    #          remaining_items.append(i[1])
    #          #  remaining_items.put((-i[1], i[1]))
    #          done_types.append(i[0])
    #  print('remaining_items', remaining_items)

    ridx = 0
    for idx, s in enumerate(S[::-1]):
        if len(remaining_items) < ridx + 1:
            break
        if len(remaining_items) == 0:
            break
        #  if remaining_items.empty():
        #      break
        if type_count_dict[s[0]] >= 2:
            # print('s', s)
            type_count_dict[s[0]] -= 1
            base_point += remaining_items[ridx] - s[1]
            #  base_point += remaining_items.get()[1] - s[1]
            type_bonus += 1
            ridx += 1
            #  results.append(type_bonus**2 + base_point)
            result = max(type_bonus**2 + base_point, result)

    #  print(max(results))
    print(result)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
ri = lambda: int(input())
rl = lambda: list(map(int,input().split()))

N,K=rl()

td_1 = [0]*N
td_rem = []
for _ in range(N):
    t,d = rl()
    t -= 1
    if td_1[t] == 0:
        td_1[t] = d
    else:
        if d <= td_1[t]:
            td_rem += [d]
        else:
            td_rem += [td_1[t]]
            td_1[t] = d


td_1.sort(reverse=True)
td_rem.sort(reverse=True)
#print(td_1)
#print(td_rem)
import numpy as np
index_0 = td_1.index(0) if 0 in td_1 else len(td_1)
td_1 = list(np.cumsum([0] + td_1[:index_0]))
td_rem = list(np.cumsum([0]+td_rem))

ans = [0 for _ in range(K+1)]
sum_l = 0
for i in range(1,K+1):
    if len(td_1)-1 >= i and len(td_rem)-1>=K-i:
        ans[i] = i*i + td_1[i] + td_rem[K-i]

#print(td_1)
#print(td_rem)
#print(ans)
print(max(ans))

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
top = [0]*n
sub = []
for i in range(n):
    t,d = map(int,input().split())
    t -=1
    if top[t]==0:
        top[t]=d
    else:
        if top[t]>=d:
            sub.append(d)
        else:
            sub.append(top[t])
            top[t]=d
 
top.sort()
top = top[::-1]
res = [0]
for i in range(n):
    if top[i]!=0:
        res.append(res[-1]+top[i])
    else:
        res.append(-1)
sub.sort()
sub = sub[::-1]
w = [0]
for i in range(len(sub)):
    w.append(w[-1]+sub[i])
tmp = 0
for i in range(1,k+1):
    if len(w)-1>=k-i and res[i]!=-1:
        tmp = max(res[i]+w[k-i]+i*i,tmp)
 
print(tmp)

Submission status: Accepted

Solution code:
ri = lambda: int(input())
rl = lambda: list(map(int,input().split()))

N,K=rl()
td = [rl() for _ in range(N)]

from operator import itemgetter
key0 = itemgetter(0)
key1 = itemgetter(1)
td.sort(key=key1,reverse=True)

netas = set()
dup = []

for t,d in td[:K]:
    if t in netas:
        dup.append(d)
    else:
        netas.add(t)

score = len(netas)**2 + sum(map(key1,td[:K]))
ans = score
for t,d in td[K:]:
    if t not in netas and dup != []:
        netas.add(t)
        dup_ = dup.pop(-1)
        score += 2*len(netas)-1 + d - dup_
        ans = max(ans,score)
print(ans)

Submission status: Accepted

Solution code:
n, k = [int(i) for i in input().split(" ")]
t = []
d = []

for i in range(0, n):
    x, y = [int(i) for i in input().split(" ")]
    t.append(x)
    d.append(y)

dt = zip(d, t)
dt = sorted(dt, key = lambda k:k[0], reverse=True)
t_max = [False for i in range(0, n + 1)]
d_max = []
d_non_max = []

for i in range(0, n):
    if t_max[dt[i][1]]:
        d_non_max.append(dt[i][0])
    else:
        t_max[dt[i][1]] = True
        d_max.append(dt[i][0])
        
p = [0 for i in range(0, len(d_max) + 1)]
sum_max = 0
sum_non_max = sum(d_non_max[0:min(len(d_non_max), k)])

for i in range(1, len(p)):
    sum_max += d_max[i - 1]
    if k - i >= 0 and k - i < len(d_non_max):
        sum_non_max -= d_non_max[k - i]
    if k - i >= 0 and k - i - 1 < len(d_non_max):
        p[i] = sum_max + sum_non_max + i * i
    else:
        p[i] = 0
        
print(max(p))

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = []
for _ in range(N):
    TD.append(list(map(int, input().split())))
TD.sort(key=lambda x: -x[1])

removable_list = []
S = set()
base_pt = 0
for t, d in TD[:K]:
    if t in S:
        removable_list.append([t, d])
    else:
        S.add(t)
    base_pt += d
kind = len(S)
removable_list.sort(key=lambda x: x[1])

addable_list = []
for t, d in TD[K:]:
    if t in S:
        continue
    addable_list.append([t, d])
    S.add(t)

pt = [base_pt + kind*kind]
l = min(len(removable_list), len(addable_list))

for rl, al in zip(removable_list[:l], addable_list[:l]):
    kind += 1
    base_pt += al[1] - rl[1]
    pt.append(base_pt + kind*kind)
print(max(pt))

Submission status: Accepted

Solution code:
import heapq
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append((-d,t))
data=sorted(data)
neta=set()
que1=[]
p,x=0,0
for i in range(k):
    if data[i][1] in neta:
        heapq.heappush(que1,(-data[i][0],data[i][1]))
        x-=data[i][0]
    else:
        x-=data[i][0]
        p+=1
        neta.add(data[i][1])
mx=x+p*p
for i in range(k,n):
    if data[i][1] in neta:
        pass
    else:
        if que1:
            a=heapq.heappop(que1)
            x=x-a[0]-data[i][0]
            p+=1
            neta.add(data[i][1])
            mx=max(mx,x+p*p)
        else:
            break
print(mx)

Submission status: Accepted

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

for x in range(0,N):
    TDlist.append(list(map(int,input().split())))

TDlist= sorted(TDlist, key = lambda x: x[1],reverse=True)
kinds = set()
seconds = []
answer = 0

for x in TDlist[:K]:
    if not x[0] in kinds:
        kinds.add(x[0])
    else:
        seconds.append(x[1])
    answer += x[1]

answer += len(kinds)**2
score = answer

for x in TDlist[K:]:
    if seconds == []:
        break
    if x[0] in kinds:
        continue
    kinds.add(x[0])
    y = seconds.pop()
    score += (x[1] - y) + (len(kinds)*2 -1)
    answer = max(score,answer)

print(answer)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
a = [[int(t) for t in input().split()] for d in range(N)]
a.sort(key=lambda x: x[1], reverse=True)
x=[]
y=[]
st = set()

for i in range(K):
    t, d = a[i]
    if(t in st):
        x.append(d)
    else:
        y.append(d)
        st.add(t)
        
X = sum(x)
Y = sum(y)
V = len(y)
res = X + Y + V ** 2

for i in range(K, N):
    if not x:
        break
    t, d = a[i]
    if t in st:
        continue
    st.add(t)
    X -= x.pop()
    Y += d
    V += 1
    res = max(res, X + Y + V ** 2)
    
print(res)