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:
from operator import itemgetter
from heapq import heappush, heappop

N, K = map(int, input().split())
a = []
for i in range(N):
    ta,tb = map(int, input().split())
    a.append([i,ta,tb])
a = sorted(a, key=itemgetter(2), reverse=True)

neta_count = 0
umai = []
umai_append = umai.append
used_i = set()
add_i = used_i.add# add is append in set()
used_neta = set()
add_neta = used_neta.add
#ネタの重複しないもののみ選んでいるつまり各ネタの種類から一種類ずつ
for i, neta, uma in a:
    if neta not in used_neta:
        umai.append(uma)
        neta_count += 1
        add_i(i)
        add_neta(neta)
    if neta_count == K:
        break
#print(umai)
m = neta_count
umai.sort()#昇順注意
for i, neta, uma in filter(lambda x: x[0] not in used_i, a):
    #print(i, neta, uma)
    if m < K:
        heappush(umai, uma)
        m += 1
    elif uma > umai[0]:
        if neta in used_neta and uma - umai[0] > neta_count**2 - (neta_count-1)**2:
            #そのネタを入れるのと種類ボーナスどっちが大きいか計算してる
            heappop(umai)
            heappush(umai, uma)
            neta_count -= 1

print(sum(umai) + neta_count**2)

Submission status: Accepted

Solution code:
from itertools import accumulate
N,K = map(int,input().split())
td = [tuple(map(int,input().split())) for i in range(N)]

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

t = [-1]*N
for i in range(N):
    t[td[i][0]] = 0
tdK = td[:K]
for i in range(K):
    t[tdK[i][0]] = 1

t_used = [i for i in range(N) if t[i]==1]
t_not_used = [i for i in range(N) if t[i]==0]

d_top = [0]*N
d_top_candi = [0]*N
d_rest = []
d_candi = []

for i in range(K):
    if d_top[td[i][0]] == 0:
        d_top[td[i][0]] = td[i][1]
    else:
        d_rest.append(td[i][1])
for i in range(K,N):
    if t[td[i][0]] == 0:
        if d_top_candi[td[i][0]] == 0:
            d_top_candi[td[i][0]] = 1
            d_candi.append(td[i][1])
l = len(t_used)
s = sum(d_top)
d_rest.sort()
d_candi.sort(key=lambda x:-x)

d_rest_cum = list(accumulate(d_rest))
d_candi_cum = list(accumulate(d_candi))

if len(d_rest) == 0:
    print(l**2 + s)
else:
    s += d_rest_cum[-1]
    ans = l**2 + s
    for k in range(min(len(d_rest),len(d_candi))):
        ans = max((l+k+1)**2 + s + d_candi_cum[k] - d_rest_cum[k],ans)

    print(ans)

Submission status: Accepted

Solution code:
from itertools import accumulate
N,K = map(int,input().split())
td = [tuple(map(int,input().split())) for i in range(N)]

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

t = [-1]*N
for i in range(N):
    t[td[i][0]] = 0
tdK = td[:K]
for i in range(K):
    t[tdK[i][0]] = 1

t_used = [i for i in range(N) if t[i]==1]
t_not_used = [i for i in range(N) if t[i]==0]

d_top = [0]*N
d_top_candi = [0]*N
d_rest = []
d_candi = []

for i in range(K):
    if d_top[td[i][0]] == 0:
        d_top[td[i][0]] = td[i][1]
    else:
        d_rest.append(td[i][1])
for i in range(K,N):
    if t[td[i][0]] == 0:
        if d_top_candi[td[i][0]] == 0:
            d_top_candi[td[i][0]] = 1
            d_candi.append(td[i][1])
l = len(t_used)
s = sum(d_top)
ans = l**2 + s
d_rest.sort()
d_candi.sort(key=lambda x:-x)

d_rest_cum = list(accumulate(d_rest))
d_candi_cum = list(accumulate(d_candi))

if len(d_rest) == 0:
    print(ans)
else:
    ans += d_rest_cum[-1]
    s += d_rest_cum[-1]
    for k in range(min(len(d_rest),len(d_candi))):
        ans = max((l+k+1)**2 + s + d_candi_cum[k] - d_rest_cum[k],ans)

    print(ans)

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-

"""
参考：https://img.atcoder.jp/abc116/editorial.pdf
"""

import sys, re
from collections import deque, defaultdict, Counter
from math import sqrt, hypot, factorial, pi, sin, cos, radians
from heapq import heappop, heappush, heapify, heappushpop
from bisect import bisect_left, bisect_right
from itertools import permutations, combinations, product
from operator import itemgetter, mul
from copy import deepcopy
from functools import reduce, partial
from fractions import Fraction
from string import ascii_lowercase, ascii_uppercase, digits

def input(): return sys.stdin.readline().strip()
def ceil(a, b=1): return int(-(-a // b))
def round(x): return int((x*2+1) // 2)
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(): return list(map(int, input().split()))
sys.setrecursionlimit(10 ** 9)
INF = float('inf')
MOD = 10 ** 9 + 7

N, K = MAP()
tdN = [[0]*2 for i in range(N)]
for i in range(N):
    t, d = MAP()
    tdN[i][0] = t
    tdN[i][1] = d

tdN.sort(key=itemgetter(1),reverse=True)
base = 0
bonus = 0
c = Counter()
changes = []
ans = 0
for i in range(K):
    t, d = tdN[i]
    base += d
    if c[t] == 0:
        bonus += 1
    c[t] += 1
    # 交換候補を集める
    if c[t] >= 2:
        heappush(changes, (d, i))
bonus *= bonus
ans = base + bonus

for i in range(K, N):
    t, d = tdN[i]
    # 既に持ってるネタならスキップ
    if c[t] != 0:
        continue
    else:
        # 1つずつ候補から交換してみて、最大値を取る
        if len(changes):
            _, idx = heappop(changes)
            t2, d2 = tdN[idx]
            base -= d2 - d
            bonus = (sqrt(bonus)+1) ** 2
            ans = max(base + bonus, ans)
            # 今回交換した分を更新
            c[t2] -= 1
            c[t] += 1
print(int(ans))

Submission status: Accepted

Solution code:
#!/usr/bin/env python3

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

    s_list.sort(key = lambda v: v[1], reverse = True)
    n_set = set()
    A, B = [0], [0]
    tot_a, tot_b = 0, 0
    for td in s_list:
        if not td[0] in n_set:
            n_set.add(td[0])
            tot_a += td[1]
            A.append(tot_a)
        else:
            tot_b += td[1]
            B.append(tot_b)

    min_n = max(1, K - (len(B) - 1))
    max_n = min(K + 1, len(A))
    max_x = 0
    for i in range(min_n, max_n):
        max_x = max(max_x, A[i] + i ** 2 + B[K - i])

    print(max_x)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
from heapq import heappop,heappush
N,K=map(int,input().split())
table=[]
for i in range(N):
    t,d=map(int,input().split())
    t=t-1
    table.append((d,t))
table=sorted(table,reverse=True)
L=[0]*N
H=set()
W=[]
ans=-1
num=0
for i in range(K):
    d,t=table[i]
    L[t]+=1
    H.add(t)
    num+=d
    if L[t]>=2:
        heappush(W,(d,t))
ans=len(H)**2+num
for i in range(K,N):
    d,t=table[i]
    L[t]+=1
    if L[t]>=2:
        continue
    if len(W)==0:
        continue
    d1,t1=heappop(W)
    H.add(t)
    num=num-d1+d
    ans=max(ans,len(H)**2+num)
print(ans)

Submission status: Accepted

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

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0
number=set()
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.add(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[index][0] not in number:
    number.add(td[index][0])
    tempanswer=tempanswer+td[index][1]-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)

Submission status: Accepted

Solution code:
def run(N, K, t, d):
    '''
    おいしさ(d)の昇順でソートする O(NlogN)
    おいしさが大きいもののK個の合計を取得 O(K)
    K個のうちおいしさが少ないものから順にネタの種類が
    増えるように変更していく O(K)
    ネタの種類のうち最大の物の値とインデックスを取得しておく O(K)
    ネタの種類を変えた合計のうち最大のものが答え
    '''
    t_dic = dict(zip(range(N), t))
    d_dic = dict(zip(range(N), d))
    d_dic_sorted = sorted(d_dic.items(), key=lambda x: x[1], reverse=True)
    d_dic_sorted_i = [dd[0] for dd in d_dic_sorted]
    d_dic_sorted_d = [dd[1] for dd in d_dic_sorted]
    d_sum = sum(d_dic_sorted_d[:K])
    t_used = {}
    t_cnt = 0
    for i in t:
        t_used[i] = 0
    for i in d_dic_sorted_i[:K]:
        t_used[t_dic[i]] += 1
        if t_used[t_dic[i]] == 1:
            t_cnt += 1
    t_sum = t_cnt ** 2
    max_sum = d_sum + t_sum
    del_list = d_dic_sorted_i[:K]
    for i in d_dic_sorted_i[K:]:
        di = d_dic[i]
        ti = t_dic[i]
        # ネタの種類が増えるように追加
        if t_used[ti] == 0:
            t_used[ti] += 1
            t_cnt += 1
            t_sum += 2*t_cnt - 1
            d_sum += di
            # ネタが重複しているものを削除
            flg = True
            while len(del_list) > 0 and flg:
                del_i = del_list.pop()
                del_di = d_dic[del_i]
                del_ti = t_dic[del_i]
                if t_used[del_ti] > 1:
                    t_used[del_ti] -= 1
                    d_sum -= del_di
                    max_sum = max(max_sum, d_sum + t_sum)
                    flg = False
    return max_sum


def main():
    N, K = map(int, input().split())
    t = []
    d = []
    for _ in range(N):
        ti, di = map(int, input().split())
        t.append(ti)
        d.append(di)
    print(run(N, K, t, d))


if __name__ == '__main__':
    main()

Submission status: Accepted

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

a = []
for i in range(N):
    t, d = map(int, input().split())
    a.append((t, d))
a.sort(key=lambda x:x[1], reverse=True)
#print(a)

vmax = 0
s = set()
q = []
for i in range(K):
    vmax += a[i][1]
    if a[i][0] in s:
        q.append(a[i][1])
    s.add(a[i][0])
vmax += len(s)**2   #add the variety bonus
#print(s)
#print(q)
#print(vmax)

i += 1   # same as i = K
v = vmax
while i < N and q:
    if a[i][0] not in s:
        v -= q.pop() + len(s)**2
        v += a[i][1] + (len(s)+1)**2
        vmax = max(vmax, v)
        s.add(a[i][0])
    i += 1
#print(s)
#print(q)
print(vmax)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x:-x[1])
 
ans = 0
s = set()
q = []
 
for i in range(k):
    ans += a[i][1]
    if a[i][0] in s:
        q.append(a[i][1])
    s.add(a[i][0])
 
ans += len(s)**2
k1 = k
ans1 = ans
 
#はじめ0~k-1までの和がansなので、k~n-1でチェック
while k1 < n and q:
    if a[k1][0] not in s:
        ans1 = ans1+a[k1][1]-q.pop()-len(s)**2+(len(s)+1)**2
        ans = max(ans, ans1)
        s.add(a[k1][0])
    k1 += 1
 
print(ans)

Submission status: Accepted

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

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


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

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

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

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

cstackd = []
nestackd = []
nestackt =[]
nneed = []
ctset = set()
sum = 0
for i in range(k):
    if a[i][1] in ctset:
        nneed.append(a[i][0])
    else : ctset.add(a[i][1])
    sum+=a[i][0]

ans = pow(len(ctset),2)+sum
for i in range(k,len(a)):
    if len(nneed) == 0: break
    if a[i][1] not in ctset: #種類がないばあい
        sum -= nneed.pop()
        sum+= a[i][0]
        ctset.add(a[i][1])
        ans = max (pow(len(ctset),2) + sum,ans)


print(ans)

Submission status: Accepted

Solution code:
n, k =map(int, input().split())
dt = sorted(tuple(map(int, input().split()))[::-1] for _ in range(n))[::-1]
r = v = 0
a = []
s = set()
for i, (d,t) in enumerate(dt, 1):
    if t in s:
        if i <= k:
            a.append(d)
            v += d
    else:
        if i > k:
            if not a:
                break
            v -= a.pop()
        v += d
        s.add(t)
    r = max(r, len(s)**2+v)
print(r)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x:-x[1])

ans = 0
s = set()
q = []

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

ans += len(s)**2
k1 = k
ans1 = ans

#はじめ0~k-1までの和がansなので、k~n-1でチェック
while k1 < n and q:
    if a[k1][0] not in s:
        ans1 = ans1+a[k1][1]-q.pop()-len(s)**2+(len(s)+1)**2
        ans = max(ans, ans1)
        s.add(a[k1][0])
    k1 += 1

print(ans)

Submission status: Accepted

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

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

ans=0
eat=set([])
noteat=[]
for t,d in td[:K]:
  ans+=d
  if t not in eat:
    eat.add(t)
  else:
    noteat.append(d)
temp_result=ans+len(eat)**2
max_result=temp_result

for t,d in td[K:]:
  if not noteat:
    break
  if t not in eat:
    temp_result+=2*len(eat)+1
    temp_result+=d-noteat.pop()
    eat.add(t)
    max_result=max(max_result,temp_result)
print(max_result)

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k != 0:
        if a[A[i][0]] ==0:
            a[A[i][0]] = 1
            eat.append(A[i][1])
            k -= 1
        else:
            noteat.append(A[i][1])
    else:
        noteat.append(A[i][1])
        
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat[k:])+1)):
    ans[i] = ans[i-1]-(pow(K-k-(i-1),2)+eat[-i]) +(pow(K-k-i,2)+noteat[k+i-1])

print(max(ans))

Submission status: Accepted

Solution code:
from collections import deque as dq
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)
ls = set()
dl = dq()
ans = 0
for i in range(K):
    ans += td[i][1]
    if not td[i][0] in ls:
        ls.add(td[i][0])
    else:
        dl.append(td[i][1])
ans = ans + len(ls)**2
res = ans
for i in range(K,N):
    if not dl:
        break
    if td[i][0] in ls:
        continue
    ans -= dl[-1]
    ans -= len(ls)*len(ls)
    ls.add(td[i][0])
    ans += len(ls)*len(ls)
    ans += td[i][1]
    del dl[-1]
    res = max(res,ans)
print(res)

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
td = [[int(i) for i in input().split()] for _ in range(N)]
td.sort(key=lambda x:x[1],reverse=True)
ans = 0
ls = set()
dl = []
for i in range(K):
    ans += td[i][1]
    if not td[i][0] in ls:
        ls.add(td[i][0])
    else:
        dl.append(td[i][1])
ans = ans + len(ls)**2
res = ans
for i in range(K,N):
    if not dl:
        break
    if td[i][0] in ls:
        continue
    ans += len(ls)*2+1
    ans += td[i][1]
    ans -= dl[-1]
    del dl[-1]
    res = max(res,ans)
    ls.add(td[i][0])
print(res)

Submission status: Accepted

Solution code:
import heapq as hq
from collections import deque
from copy import deepcopy
N, K = map(int, input().split())
Sushi = [None for _ in range(N)]
TotalType = set()
for _ in range(N):
    t, d = map(int, input().split())
    Sushi[_] = [d, t]
Sushi.sort(reverse = True)

ChosenType = set()
Double = []
hq.heapify(Double)
BasePoint = 0
for goodsushi in Sushi[:K]:
    BasePoint += goodsushi[0]
    if goodsushi[1] in ChosenType:
        hq.heappush(Double, goodsushi)
    else: ChosenType |= {goodsushi[1]}

NotChosen = deque(deepcopy(Sushi[K:]))
MaxHappiness = BasePoint + len(ChosenType)**2
while Double and NotChosen:
    TrushPoint, TrushType = hq.heappop(Double)
    BasePoint -= TrushPoint
    while NotChosen:
        addPoint, addType = NotChosen.popleft()
        if addType not in ChosenType:
            BasePoint += addPoint
            ChosenType |= {addType}
            break
    MaxHappiness = max(MaxHappiness, BasePoint + len(ChosenType)**2)
print(MaxHappiness)

Submission status: Accepted

Solution code:
from collections import *
n, k = map(int, input().split())
d, t = zip(*sorted(tuple(map(int, input().split()))[::-1] for _ in range(n))[::-1])
c = Counter(t[:k])
x = len(c)
a = [sum(d[:k])]
for i in range(k)[::-1]:
    if c[t[i]] > 1:
        c[t[i]] -= 1
        a.append(a[-1] - d[i])
b = [0]
for i in range(k, n):
    if c[t[i]] <= 0:
        c[t[i]] += 1
        b.append(b[-1] + d[i])
r = 0
for i, (v, w) in enumerate(zip(a, b)):
    r = max(r, v+w+(x+i)**2)
print(r)

Submission status: Accepted

Solution code:
import sys

def r(n, k, sl):

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

    gmk = set()
    gop = 0
    ce = []
    for ti, di in sl[:k]:
        if ti in gmk:
            ce.append(di)
        gmk.add(ti)
        gop += di

    gscr = len(gmk)**2 + gop

    for ti, di in sl[k:]:
        if 0 >= len(ce):
            break
        if ti in gmk:
            continue
        gmk.add(ti)
        odi = ce.pop()
        gop = gop - odi + di
        nscr = len(gmk)**2 + gop
        gscr = max(gscr, nscr)

    return gscr


def main():
    input = sys.stdin.readline
    n, k = list(map(int, input().split()))
    sl = [tuple(map(int, input().split())) for _ in range(n)]
    print(r(n, k, sl))


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
N,K = map(int,input().split())
td = [[int(i) for i in input().split()] for _ in range(N)]
td.sort(key=lambda x:x[1],reverse=True)
ans = 0
ls = set()
dl = []
for i in range(K):
    ans += td[i][1]
    if not td[i][0] in ls:
        ls.add(td[i][0])
    else:
        dl.append(td[i][1])
ans = ans + len(ls)**2
res = ans
for i in range(K,N):
    if not dl:
        break
    if td[i][0] in ls:
        continue
    ans += len(ls)*2+1
    ans += td[i][1]
    ans -= dl[-1]
    del dl[-1]
    res = max(res,ans)
    ls.add(td[i][0])
print(res)

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-


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

# Store data
sushis = []
neta_set = set()
for i in range(n):
    t, d = map(int, input().split())
    sushis.append((t,d))
    neta_set.add(t)


# Formatting
sushis.sort(key=lambda i:i[1])
tMax = len(neta_set)


# Select sushi set as maximum oishisa point
sel_sushis = []
neta_set = set()
sum_oishisa = 0
for i in range(k):
    t, d = sushis.pop()
    if t in neta_set:
        sel_sushis.append((t,d))
    else:
        neta_set.add(t)
    sum_oishisa += d

sel_sushis.sort(key=lambda i:i[1], reverse=True)
tMin = len(neta_set)
max_score = tMin**2 + sum_oishisa

for nNeta in range(tMin+1, tMax + 1, 1):

    # Find sushi having minimum oishisa and being not unique for neta
    if len(sel_sushis) == 0:
        break
    t1, d1 = sel_sushis.pop()
    sum_oishisa -= d1

    # Find sushi having maximum oishisa and unique for neta
    while True:
        t2, d2 = sushis.pop()
        if not t2 in neta_set:
            break

    sum_oishisa += d2
    neta_set.add(t2)

    score = nNeta**2 + sum_oishisa
    if score > max_score:
        max_score = score

print(max_score)

Submission status: Accepted

Solution code:
from math import pow
from collections import deque, Counter
import sys
stdin = sys.stdin
def li(): return map(int, stdin.readline().split())

def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    compare_queue = deque(sorted_recs[K:])
    distinct_set = set()
    dup_queue = deque()
    score = 0
    for neta_point in target_queue:
        score += neta_point[1]
        if neta_point[0] not in distinct_set:
            distinct_set.add(neta_point[0])
        else:
            dup_queue.append(neta_point)

    cardinality = len(distinct_set)
    score += pow(cardinality, 2)
    score = int(score)
    maybe_score = score
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in distinct_set:
            continue
        elif len(dup_queue) < 1:
            break
        else:
            old = dup_queue.pop()
            inc = 2 * len(distinct_set) + 1
            dec = old[1] - comp[1]
            delta = inc - dec
            maybe_score = maybe_score + delta
            score = max(score, maybe_score)

            distinct_set.add(comp[0])
    # both queues are drained
    return score


N, K = li()
records = [tuple(li()) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Accepted

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

candidate = []
candidateset = set()
donyokulength = len(donset)#len(donset)の長さが変わってしまうので予めとっておく
for i in range(k, n):
    if f[i][0] not in donset:#
        candidate.append(-f[i][1])#最小のものから取り出されるのでマイナスをつけることによって最大のものから取り出せるようにしている
        candidateset.add(f[i][0])
        donset.add(f[i][0])#donsetにも入れて二番目以降の侵入を妨げる

satili = [donyoku]
donyoku -= donyokulength**2
for i in range(len(candidateset)):
    if notfirst == []:
        break
    else:
        donyoku -= heappop(notfirst)
        donyoku += -heappop(candidate)#最小のものから取り出されるのでマイナスをつけることによって最大のものから取り出せるようにしている
        donyoku += (donyokulength + i + 1) ** 2
        satili.append(donyoku)
    donyoku -= (donyokulength + i + 1) ** 2




print(max(satili))

Submission status: Accepted

Solution code:
from heapq import heapify,heappush,heappop
N,K=map(int,input().split())
sushi=[]
for i in range(N):
  t,d=map(int,input().split())
  sushi.append((d,t-1))
sushi.sort(key=lambda x:-x[0])
inf=float("inf")
oishisa=[-inf]*(N+1)
neta=[0]*N
heaps=[[] for i in range(N)]
S=[]
heapify(S)
for i in range(N):
  heapify(heaps[i])
a,k=0,0
for d,t in sushi[:K]:
  a+=d
  neta[t]+=1
  if neta[t]==1:
    k+=1
  heappush(heaps[t],d)

oishisa[k]=a
for i in range(N):
  if neta[i]>=2:
    d=heappop(heaps[i])
    heappush(S,(d,i))
for d,t in sushi[K:]:
  if len(S)==0:
    break
  if neta[t]==0:
    neta[t]+=1
    d_,t_=heappop(S)
    neta[t_]-=1
    if neta[t_]>=2:
      d__=heappop(heaps[t_])
      heappush(S,(d__,t_))
    heappush(heaps[t],d)
    oishisa[k+1]=oishisa[k]+d-d_
    k+=1
ans=0
for i in range(N+1):
  ans=max(oishisa[i]+i*i,ans)
print(ans)

Submission status: Accepted

Solution code:
from math import pow
from collections import deque, Counter
import sys
stdin = sys.stdin
def li(): return map(int, stdin.readline().split())

def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    compare_queue = deque(sorted_recs[K:])
    distinct_set = set()
    dup_queue = deque()
    score = 0
    for neta_point in target_queue:
        score += neta_point[1]
        if neta_point[0] not in distinct_set:
            distinct_set.add(neta_point[0])
        else:
            dup_queue.append(neta_point)

    cardinality = len(distinct_set)
    score += pow(cardinality, 2)
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in distinct_set:
            continue
        elif len(dup_queue) < 1:
            break
        else:
            old = dup_queue.pop()
            inc = 2 * len(distinct_set) + 1
            dec = old[1] - comp[1]
            delta = inc - dec
            new_score = new_scores[-1] + delta
            new_scores.append(new_score)

            distinct_set.add(comp[0])
    # both queues are drained
    return max(new_scores)

N, K = li()
records = [tuple(li()) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Accepted

Solution code:
import sys
stdin = sys.stdin
 
sys.setrecursionlimit(10**5) 
 
def li(): return map(int, stdin.readline().split())
def li_(): return map(lambda x: int(x)-1, stdin.readline().split())
def lf(): return map(float, stdin.readline().split())
def ls(): return stdin.readline().split()
def ns(): return stdin.readline().rstrip()
def lc(): return list(ns())
def ni(): return int(stdin.readline())
def nf(): return float(stdin.readline())

n,k = li()
td = [tuple(li()) for _ in range(n)]

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

kindset = set()
doublestack = []

ans = 0

for ti, di in td[:k]:
    if ti in kindset:
        doublestack.append((ti,di))
    
    else:
        kindset.add(ti)
        
    ans += di

ans += len(kindset) * len(kindset)
score = ans

# n-k個の要素に対して重複を排して交換
for ti, di in td[k:]:
    if ti in kindset:
        continue
    
    # すべてユニークだったら抜ける
    if len(doublestack) == 0:
        break
    
    #スコア再計算
    old_ti, old_di = doublestack.pop()    
    score = score - old_di + di + 2*len(kindset) + 1
    
    # 更新
    ans = max(ans, score)    
    kindset.add(ti)
    
    
print(ans)

Submission status: Accepted

Solution code:
N, K = [int(x) for x in input().split()]

# sushi
s = {}
for i in range(N):
    t, d = [int(x) for x in input().split()]
    if not t in s:
        s[t] = [d]
    else:
        s[t].append(d)

# 各スシのMax
m = [max(v) for k, v in s.items()]
m.sort(reverse=True)

s2 = {}
all2 = []
for k in s.keys():
    s2[k] = sorted(s[k])
    s2[k].pop()
    all2 += s2[k]
all2.sort(reverse=True)

# スシ種類
# kinds = len(set([k for k in s.keys()]))
kinds = len(m)

import itertools

acM = list(itertools.accumulate(m))
acAll2 = list(itertools.accumulate(all2))

ans = 0
# for i in range(1, min(K, kinds) + 1):
for i in range(0, min(K, kinds)):
    # addM = m[:i]
    # addAll2 = all2[:K - i]
    # newAns = sum(addM) + sum(addAll2) + i ** 2
    t1 = acM[i]
    idx = K - i - 2
    t2 = acAll2[idx] if idx >= 0 and idx < len(acAll2)  else 0
    newAns = t1 + t2 + (i + 1) ** 2
    ans = max(ans, newAns)

print(ans)

pass

Submission status: Accepted

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


n, k = map(int, input().split())
dic = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    dic[t].append(d)

maxs = []
for key in dic:
    dic[key].sort()
    maxs.append((dic[key].pop(), key))

maxs.sort(reverse=True)
queSize = 0
que = []
queScore = 0
defaultScore = 0
cnt = 0
ans = 0
for v, key in maxs:
    defaultScore += v + cnt * 2 + 1
    cnt += 1
    if cnt >= k:
        ans = max(ans, defaultScore)
        break
    for add in dic[key]:
        heappush(que, add)
        queSize += 1
        queScore += add
    while queSize > k - cnt:
        queScore -= heappop(que)
        queSize -= 1
    ans = max(ans, defaultScore + queScore)

print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict as dd
import heapq

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

sushi_types = set()
sushies = []
v=0

for i in range(k):
    t,d=l[i]
    v+=d
    if t in sushi_types:
        heapq.heappush(sushies,d)
    else:
        sushi_types.add(t)

remaining_sushi=dd(list)
for i in range(k,n):
    t,d=l[i]
    if not t in sushi_types:
        heapq.heappush(remaining_sushi[t],-d)

eat_candidate_sushi = []
for i in remaining_sushi.keys():
    if not i in sushi_types:
        heapq.heappush(eat_candidate_sushi,remaining_sushi[i][0])

sushi_kinds=len(sushi_types)
val=v
ans=v+sushi_kinds**2

for i in range(sushi_kinds+1,k+1):
    if sushies and eat_candidate_sushi:
        v-=heapq.heappop(sushies)
        v-=heapq.heappop(eat_candidate_sushi)
        ans=max(ans,v+i**2)
    else:
        break

print(ans)

Submission status: Accepted

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

inputs = []
priority_queue = []
queue = []
sets = set()

for _ in range(N):
    key, value = map(int, input().split())
    inputs.append((value, key))

inputs.sort(reverse=True) # valueで降順のソートをかけてる

for i in range(K):
    value, key = inputs[i]
    if key in sets:
        queue.append(value)
    else:
        priority_queue.append(value)
        sets.add(key)

sum_priority_queue = sum(priority_queue)
sum_queue = sum(queue)
kinds = len(sets)
ans = sum_priority_queue + sum_queue + kinds ** 2
for i in range(K, N):
    if not queue:
        break
    value, key = inputs[i]
    if key in sets:
        continue
    sets.add(key)
    sum_queue -= queue.pop()
    sum_priority_queue += value
    kinds += 1
    ans = max(ans, sum_priority_queue + sum_queue + kinds ** 2)

print(ans)

Submission status: Accepted

Solution code:
# 解説放送

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

e = [None for _ in range(N + 1)] # 種別の美味しさ
t_kinds_count = 0 # 何種類あるか
rest_length = 0 # 各種の2番手以降の個数の総数

tops = [] # 各種の最大美味しさ
c_tops = [0] # 各種の最大美味しさの累積和
rests = [] # 各種の2番手以降
c_rests = [0] # 各種の2番手以降の累積和

for _ in range(N):
    t, d = map(int, input().split())
    if e[t] is None:
        e[t] = []
        t_kinds_count += 1
    e[t].append(d)


for t in range(1, N + 1):
    if e[t] is None:
        continue
    e[t].sort(reverse=True)
    tops.append(e[t][0])
    rests.extend(e[t][1:])

tops.sort(reverse=True)
rests.sort(reverse=True)

s = 0
for tt in tops:
    s += tt
    c_tops.append(s)

s = 0
for rr in rests:
    s += rr
    c_rests.append(s)
    rest_length += 1

ans = 0
for yy in range(1, min(t_kinds_count, K) + 1):
    if 0 <= K - yy <= rest_length:
        ans = max(ans, c_tops[yy] + c_rests[K - yy] + yy * yy)

# print(e)
# print(tops)
# print(c_tops)
# print(rests)
# print(c_rests)
print(ans)