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:
#!/usr/bin/env python3
import sys
try:
    from typing import List
except ImportError:
    pass


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    ss = sorted(zip(d, t), reverse=True)
    gs = []
    mst = set()
    for di, ti in ss[:K]:
        if ti in mst:
            gs.append((di, ti))
        else:
            mst.add(ti)
    dm = {}
    for di, ti in ss[K:]:
        if ti in dm or ti in mst:
            continue
        dm[ti] = di
    hs = sorted(((di, ti) for ti, di in dm.items()), reverse=True)
    cd = sum(di for di, ti in ss[:K])
    ct = len(mst)
    m = cd + ct * ct
    for (gdi, gti), (hdi, hti) in zip(reversed(gs), hs):
        cd += hdi - gdi
        ct += 1
        m = max(m, cd + ct * ct)
    print(m)


def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]"
    d = [int()] * (N)  # type: "List[int]"
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
N,K,*L = map(int, open(0).read().split())
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  ls.append((d,t))
ls.sort()
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = ls.pop()
  ans += d
  if t not in S:
    num += 1
    S.add(t)
  else:
    pre.append(d)
pre.sort(reverse=True)
ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = ls.pop()
  if t in S:
    continue
  S.add(t)
  m -=pre.pop()
  m += d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict
from heapq import heappush, heappop
N,K,*L = map(int, open(0).read().split())
dic = defaultdict(list)
ls = []
for i,(t,d) in enumerate(zip(*[iter(L)]*2)):
  dic[t].append(d)
  heappush(ls,(-d,t))
S = set()
pre = []
ans = 0
num = 0
for i in range(K):
  d,t = heappop(ls)
  ans -= d
  if t not in S:
    num += 1
    S.add(t)
  else:
    heappush(pre,-d)

ans += num*num
if pre==[]:
  print(ans)
  exit()
m = ans
for i in range(N-K):
  d,t = heappop(ls)
  if t in S:
    continue
  S.add(t)
  m -= heappop(pre)
  m += -d+2*num+1
  num += 1
  ans = max(ans,m)
  if pre==[]:
    break
print(ans)

Submission status: Accepted

Solution code:
import heapq
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)

NETA=set()
duplication=[]
heapq.heapify(duplication)
ans=0
for t,d in TD[:K]:
    if t in NETA:
        heapq.heappush(duplication,d)
    else:
        NETA.add(t)
    ans+=d

s=[-float("inf")]*(K+1)
cnt=len(NETA)
s[cnt]=ans
for t,d in TD[K:]:
    if t in NETA:
        continue
    else:
        if duplication:
            x=heapq.heappop(duplication)
            ans+=d-x
            NETA.add(t)
            cnt+=1
            s[cnt]=ans

for i in range(K+1):
    s[i]+=i*i
print(max(s))

Submission status: Accepted

Solution code:
import heapq
n,k=map(int,input().split())
td=[[0]*2 for _ in range(n)]
for i in range(n):
    t,d=map(int,input().split())
    td[i][0],td[i][1]=d,t
td.sort(reverse=True)
di=set()
sa=[]
heapq.heapify(sa)
ans=0
for d,t in td[:k]:
    if t in di:
        heapq.heappush(sa, d)
    else:
        di.add(t)
    ans+=d

s=[-10**18]*(k+1)
cnt=len(di)
s[cnt]=ans
for d,t in td[k:]:
    if t in di:
        continue
    else:
        if sa:
            x=heapq.heappop(sa)
            ans+=d-x
            di.add(t)
            cnt+=1
            s[cnt]=ans

for i in range(k+1):
    s[i]+=i*i

print(max(s))

Submission status: Accepted

Solution code:
from operator import itemgetter
n,k=map(int,input().split())
td=[list(map(int,input().split())) for _ in range(n)]
td.sort(key=itemgetter(1),reverse=True)
ans=0
kind=0
flg=[False]*n
sing,comp=[],[]
for i in range(n):
    t,d=td[i]
    t-=1
    if i<k:
        ans+=d
        if flg[t]:
            comp.append(d)
        else:
            flg[t]=True
            kind+=1
    else:
        if not flg[t]:
            sing.append(d)
            flg[t]=True
comp=comp[::-1]
cnt=ans
ans+=kind**2
#print(comp)
#print(sing)
if len(comp)>len(sing):
    for i in range(len(sing)):
        dc,ds=comp[i],sing[i]
        cnt+=ds-dc
        kind+=1
        ans=max(ans,cnt+kind**2)
else:
    for i in range(len(comp)):
        dc,ds=comp[i],sing[i]
        cnt+=ds-dc
        kind+=1
        ans=max(ans,cnt+kind**2)
print(ans)

Submission status: Accepted

Solution code:
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)
cand = []
memo = [0 for i in range(n+1)]

ans = 0
duplicate = []
unique = []
num = 0
i = 0

for t, d in td:
    i += 1
    if i <= k:
        ans += d
        if memo[t]:
            duplicate.append(d)
        else:
            memo[t] = 1
            num += 1
    else:
        if not memo[t]:
            unique.append(d)
            memo[t] = 1

tmp = ans
ans += num * num

# print(tmp, ans)

for d, u in zip(reversed(duplicate), unique):
    tmp += u - d
    num += 1
    ans = max(ans, tmp + num * num)

print(ans)

Submission status: Accepted

Solution code:
def main():
    n, k = map(int, input().split())
    sushi = [list(map(int, input().split())) for _ in range(n)]
    sushi.sort(key=lambda x: x[1], reverse=True)
    ans = 0
    t = [0]*(n+1)
    tmp = 0
    cnt = 0
    for i in range(k):
        ans += sushi[i][1]
        if t[sushi[i][0]] == 0:
            cnt += 1
        tmp += sushi[i][1]
        t[sushi[i][0]] += 1
    ans = tmp + cnt*cnt
    p = k
    for i in range(k, n):
        f = False
        if t[sushi[i][0]] == 0:
            f = True
            for j in reversed(range(p)):
                if 1 < t[sushi[j][0]]:
                    t[sushi[j][0]] -= 1
                    t[sushi[i][0]] += 1
                    tmp -= sushi[j][1]
                    tmp += sushi[i][1]
                    p = j
                    cnt += 1
                    ans = max(ans, tmp+cnt*cnt)
                    f = False
                    break
        if f:
            break
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
from collections import *
from heapq import *

N, K = map(int, input().split())
td = [tuple(map(int, input().split())) for _ in range(N)]
td.sort(key=lambda t: t[1], reverse=True)
d = defaultdict(list)
now = 0

for ti, di in td[:K]:
    d[ti].append(di)
    now += di

remove = []

for k in d.keys():
    d[k].sort(reverse=True)
    
    for i in range(1, len(d[k])):
        heappush(remove, d[k][i])
    
s = set(list(d.keys()))
add = []

for ti, di in td:
    if ti in s:
        continue
    
    s.add(ti)
    heappush(add, -di)

x = len(d)
ans = now+x**2

while len(remove)>0 and len(add)>0:
    now -= heappop(remove)
    now -= heappop(add)
    x += 1
    ans = max(ans, now+x**2)

print(ans)

Submission status: Accepted

Solution code:
# input()
# int(input())
# map(int, input().split())
# list(map(int, input().split()))

inf = 10**9+7
mod = 10**9+7
n, k = map(int, input().split())
g_val = [list(map(int, input().split())) for i in range(n)]
netalist = [False for i in range(n+1)]
netanum = [0 for i in range(n+1)]
g_val.sort(key=lambda x: x[1])
used_list = []
ans = [0 for i in range(k+1)] # i種類食べた時の最大値
preans = 0
for j in range(k):
    if g_val:
        g,v = g_val.pop()
        netalist[g] = True
        netanum[g]+=1
        used_list.append((v,g))
        preans+=v

kind = sum(list(map(int, netalist)))
ans[kind] = preans+(kind*kind)

import heapq # 優先度付きキュー(最小値取り出し)
heapq.heapify(used_list)   # ヒープ化


while used_list:
    v,g = heapq.heappop(used_list) # 最小値の取り出し
    if netanum[g] > 1:
        while g_val:
            gnew,vnew = g_val.pop()
            if netanum[gnew] == 0:
                ans[kind+1] = ans[kind]-v-(kind*kind)+vnew+((kind+1)*(kind+1))
                kind+=1
                netanum[g] -= 1
                netanum[gnew] += 1
                break
    if not g_val:
        break
print(max(ans))

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
s = set()
l = []
ln = [[] for i in range(n)]
for i in range(n):
    t,m = map(int,input().split())
    l.append([m,t-1])
    ln[t-1].append(m)
l.sort(reverse = True)
cand = []

ans = 0
mod = []
for i in range(k):
    ans += l[i][0]
    if l[i][1] in s:
        mod.append(l[i][0])
    else:
        s.add(l[i][1])
x = len(s)
now = ans + x**2

for i in range(n):
    
    if i in s or ln[i] == []:
        continue
    cand.append(max(ln[i]))

cand.sort(reverse = True)

for i in range(min(len(mod),len(cand))):
    ans += cand[i]-mod[-1-i]
    x += 1
    now = max(now,ans+x**2)
print(now)

Submission status: Accepted

Solution code:
# coding: utf-8
import sys
from heapq import heapify, heappop, heappush

sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))

# x * xをどう評価するか、heapに「おいしさ」の高い順に入れていく
N, K = lr()
TD = [lr() for _ in range(N)]
heap = []
kind = 0
happy = 0
used = set()
TD.sort(key=lambda x: x[1], reverse=True)
for t, d in TD[:K]:
    if t not in used:
        used.add(t)
        kind += 1
    else:
        heappush(heap, d)
    happy += d

cand = [happy+kind**2]
for t, d in TD[K:]:
    if not heap:
        break
    if t in used:
        continue
    h = heappop(heap)
    happy += (d-h)
    kind += 1
    used.add(t)
    cand.append(happy+kind**2)

answer = max(cand)
print(answer)

Submission status: Accepted

Solution code:
from collections import defaultdict
N,K=map(int,input().split())

dtlist=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tset_all.add(t)
dtlist.sort()
#print(heapq)

dtlist_Krev=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  d,t=dtlist.pop()
  dtlist_Krev.append((d,t))
  dsum+=d
  tdic[t]+=1
#dtlist_Krev.sort(reverse=True)

t0=len(tdic)
max_answer=dsum+(t0**2)
#print(t0,max_answer)

for i in range(t0+1,min(K,len(tset_all))+1):
  while(True):    
    d1,t1=dtlist_Krev.pop()
    if tdic[t1]>1:
      tdic[t1]-=1
      dsum-=d1
      break
      
  while(True):        
    d2,t2=dtlist.pop()
    if tdic[t2]==0:
      tdic[t2]=1
      dsum+=d2      
      break
          
  answer_i=dsum+(i**2)
  #print(i,answer_i)
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

Submission status: Accepted

Solution code:
from collections import defaultdict
N,K=map(int,input().split())

dtlist=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  dtlist.append((d,t))
  tset_all.add(t)
dtlist.sort()
#print(heapq)

dtlist_Krev=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  d,t=dtlist.pop()
  dtlist_Krev.append((d,t))
  dsum+=d
  tdic[t]+=1
dtlist_Krev.sort(reverse=True)

t0=len(tdic)
max_answer=dsum+(t0**2)
#print(t0,max_answer)

for i in range(t0+1,min(K,len(tset_all))+1):
  loop_flg=True
  while(loop_flg):    
    d,t=dtlist_Krev.pop()
    if tdic[t]>1:
      tdic[t]-=1
      dsum-=d
      
      while(True):        
        d2,t2=dtlist.pop()
        if tdic[t2]==0:
          tdic[t2]=1
          dsum+=d2
          loop_flg=False
          break
          
  answer_i=dsum+(i**2)
  #print(i,answer_i)
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

Submission status: Accepted

Solution code:
from collections import defaultdict
import heapq
N,K=map(int,input().split())

hq=[]
tset_all=set()
for i in range(N):
  t,d=map(int,input().split())
  heapq.heappush(hq,(-d,t))
  tset_all.add(t)
#print(heapq)

hq_K=[]
dsum=0
tdic=defaultdict(int)
for i in range(K):
  md,t=heapq.heappop(hq)
  heapq.heappush(hq_K,((-md,t)))
  dsum-=md
  tdic[t]+=1  

t0=len(tdic)
max_answer=dsum+(t0**2)
#print(t0,max_answer)

for i in range(t0+1,min(K,len(tset_all))+1):
  loop_flg=True
  while(loop_flg):
    d,t=heapq.heappop(hq_K)
    if tdic[t]>1:
      tdic[t]-=1
      dsum-=d
      
      while(True):
        md2,t2=heapq.heappop(hq)
        if tdic[t2]==0:
          tdic[t2]=1
          dsum-=md2
          loop_flg=False
          break
          
  answer_i=dsum+(i**2)
  #print(i,answer_i)
  max_answer=max(max_answer,answer_i)
  
print(max_answer)

Submission status: Accepted

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

n, k = map(int, input().split())
TD = [[0, 0] for _ in range(n)]
for i in range(n):
    t, d = map(int, input().split())
    TD[i][0], TD[i][1] = d, t

TD.sort(reverse=True)
S = set()
f = 0
h = []
import heapq
heapq.heapify(h)
for i in range(k):
    d, t = TD[i]
    if t in S:
        heapq.heappush(h, d)
    else:
        S.add(t)
    f += d

INF = 10**18
F = [-INF]*(k+1)
cur = len(S)
F[cur] = f

for i in range(k, n):
    d, t = TD[i]
    if t in S:
        continue
    else:
        if h:
            x = heapq.heappop(h)
            f += d-x
            S.add(t)
            cur += 1
            F[cur] = f

for i in range(0, k+1):
    F[i] += i*i

print(max(F))

Submission status: Accepted

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

from collections import defaultdict

def main():
    n, k = map(int, input().split())
    td = []
    for i in range(n):
        ti, di = map(int, input().split())
        td.append((ti, di))
    td = list(reversed(sorted(td)))
    tdd = defaultdict(list)
    for t, d in td:
        tdd[t].append(d)

    a = []; b = []
    for _, v in tdd.items():
        a.append(v[0])
        b += v[1:]

    a = list(reversed(sorted(a)))
    b = list(reversed(sorted(b)))
    for i in range(1, len(a)):
        a[i] += a[i - 1]
    for i in range(len(a)):
        a[i] += (i + 1) ** 2
    for i in range(1, len(b)):
        b[i] += b[i - 1]

    a = [0] + a
    b = [0] + b

    res = 0
    for i in range(1, k + 1):
        if i >= len(a) or (k - i) >= len(b):
            continue
        ts = a[i] + b[k - i]
        res = max(res, ts)

    print(res)

if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
# solution
# online help seeked

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key=lambda x: x[1])
t = {}
x = []
z = 0
for i in range(n - 1, n - k - 1, -1):
    ti, di = sushi[i]
    if ti in t.keys():
        x.append(di)
        t[ti] += 1
    else:
        t[ti] = 1
    z += di

p = len(t.keys())
ans = z + p ** 2
idx = 1
for i in range(n - k - 1, -1, -1):
    if idx == len(x) + 1:
        break
    ti, di = sushi[i]
    if ti not in t.keys():
        t[ti] = 1
        z += di - x[-idx]
        p += 1
        idx += 1
        ans = max(ans, z + p ** 2)

print(ans)

Submission status: Accepted

Solution code:
n,k = map(int,input().split())
td = [list(map(int,input().split())) for i in range(n)]
td.sort(key=lambda x:x[1],reverse=True)
a = 0
ans = []
kind = [-1]*n
c = 0
min_p = []
for i in range(n):
    eat = False
    if i<k:
        a += td[i][1]
        eat = True
    else:
        if min_p:
            if kind[td[i][0]-1]<0:
                a -= (min_p.pop()-td[i][1])
                a += (c+1)**2-c**2
                ans.append(a)
                eat = True
    if eat:
        if kind[td[i][0]-1]<0:
            kind[td[i][0]-1]=1
            c+=1
        else:min_p.append(td[i][1])
    if i==k-1:
        a += c**2
        ans.append(a)
print(max(ans))

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],reverse=True)
td_core=[]
td_sel=[]
td_sub=[]
cat=[0]*(n+1)
for i in range(n):
  if i < k and cat[td[i][0]]==0:
    td_core.append(td[i])
    cat[td[i][0]]=1
  elif i < k:
    td_sel.append(td[i])
  elif cat[td[i][0]]==0:
    td_sub.append(td[i])
    cat[td[i][0]]=1
# i >= k and cat[td[i][0]]==1 となるものはすでにtd_core,td_sel,td_sub内ものと種類も被っている上に美味しさも低い
# 後の工程は種類を増やす処理なので、同じ種類のネタはそのネタの中で美味しさが最大のもの一つで十分。なので捨てる

# dの大きい順に貪欲にネタを取得したときの満足ポイント
kiso=sum([x[1] for x in td_core])+sum([x[1] for x in td_sel])
bns=len(td_core)
ans=kiso+bns**2
# この状態からネタの種類を増やしていき最大が出てくるか検証
# td_sel内のネタは同一ネタが複数あり、そのネタ内で最大じゃないものが美味しさの大きい順に入っている。
# td_subにはtd_core,td_selにないネタの最大美味しさの寿司が入っている。
td_sel.sort(key=lambda x:x[1])
for i in range(min(len(td_sel),len(td_sub))):
  kiso-=td_sel[i][1]
  kiso+=td_sub[i][1]
  bns+=1
  ans=max(ans,kiso+bns**2)
print(ans)

Submission status: Accepted

Solution code:
#!usr/bin/env python3
from collections import defaultdict, deque
from heapq import heappush, heappop
from itertools import permutations, accumulate
import sys
import math
import bisect
def LI(): return [int(x) for x in sys.stdin.buffer.readline().split()]
def I(): return int(sys.stdin.buffer.readline())
def LS():return [list(x) for x in sys.stdin.readline().split()]
def S():
    res = list(sys.stdin.readline())
    if res[-1] == "\n":
        return res[:-1]
    return res
def IR(n):
    return [I() for i in range(n)]
def LIR(n):
    return [LI() for i in range(n)]
def SR(n):
    return [S() for i in range(n)]
def LSR(n):
    return [LS() for i in range(n)]

sys.setrecursionlimit(1000000)
mod = 1000000007

def solve():
    n,k = LI()
    g = LIR(n)
    l = list(set(([i for (i,j) in g])))
    g.sort(key=lambda x:-x[1])
    l.sort()
    q = []
    q2 = []
    f = [0]*len(l)
    ans = 0
    res = k
    for i in range(n):
        t,d = g[i]
        g[i][0] = t = bisect.bisect_left(l,t)
        if res:
            f[t] += 1
            ans += d
            res -= 1
            if f[t] > 1:
                heappush(q,(d,t))
        else:
            heappush(q2,(-d,t))
    x = sum([i >= 1 for i in f])
    ans += x**2
    s = ans
    while q2 and q:
        d,t = heappop(q2)
        d *= -1
        if f[t]:
            continue
        else:
            nd,nt = heappop(q)
            s -= nd
            s += 2*x+1
            s += d
            x += 1
            f[t] = 1
            f[nt] -= 1
            if ans < s:
                ans = s
    print(ans)
    return

#Solve
if __name__ == "__main__":
    solve()

Submission status: Accepted

Solution code:
import sys
from heapq import heapify, heappop, heappush

def solve():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    S = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        S[i] = (-d, t)
    heapify(S)

    base_point = 0
    appeared_type = set()
    type_Num = 0
    duplicated = []
    heapify(duplicated)
    for i in range(K):
        d, t = heappop(S)
        if t in appeared_type:
            base_point -= d
            heappush(duplicated, -d)
        else:
            base_point -= d
            appeared_type |= {t}
            type_Num += 1

    Ans = base_point + type_Num * type_Num
    while S:
        d, t = heappop(S)
        if t in appeared_type: continue
        else:
            if not duplicated: continue
            else:
                remove = heappop(duplicated)
                base_point -= d + remove
                type_Num += 1
                appeared_type |= {t}
                Ans = max(Ans, base_point + type_Num * type_Num)
    print(Ans)

    return 0

if __name__ == "__main__":
    solve()

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
td=[list(map(int,input().split()))[::-1]for _ in range(n)]
td.sort(reverse=1)
s=set()
ss=set()
a=[]
b=[]
c=[]
for i in range(n):
  if i<k:
    d,t=td[i]
    if t in s:
      b.append(d)
    else:
      a.append(d)
      s.add(t)
      ss.add(t)
  else:
    d,t=td[i]
    if t not in ss:
      c.append(d)
      ss.add(t)
c=c[::-1]
ls=len(s)
su=sum(a)+sum(b)
ans=ls*ls+su
for i in range(min(len(b),len(c))):
  ls+=1
  su-=b.pop()
  su+=c.pop()
  ans=max(ans,ls*ls+su)
print(ans)

Submission status: Accepted

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

ans=0
kl=dict()
for i in range(k):
    ans+=td[i][1]
    if td[i][0] in kl:
        kl[td[i][0]]+=1
    else:
        kl[td[i][0]]=1
l=len(kl)
ans+=(l**2)
ans_=ans
now=k-1
for i in range(k,n):
    if td[i][0] not in kl:
        #最小を求めるのをもっと早く
        while now>=0:
            if kl[td[now][0]]>1:
                mi=td[now]
                kl[mi[0]]-=1
                now-=1
                break
            now-=1
        if now==-1:
            break
        else:
            ans=ans+2*l+1-mi[1]+td[i][1]
            kl[td[i][0]]=1
            ans_=max(ans,ans_)
            l+=1

print(ans_)

Submission status: Accepted

Solution code:
def Int():
    return int(input())
def Ints():
    return map(int,input().split())
def IntList():
    return list(Ints())
def IntMat(N):
    return [IntList() for i in range(N)]

import sys
sys.setrecursionlimit(4100000)
input = sys.stdin.readline

N,K = Ints()
ALL = []
NetaMax = 0
MaxNeta = {}

for i in range(N):
    t,d = Ints()
    ALL.append([d,t])
    if not t in MaxNeta:
        NetaMax += 1
        MaxNeta[t] = d
    else:
        MaxNeta[t] = max(MaxNeta[t], d)
    
ALL.sort()
ALL = ALL[::-1]
NetaBest = []

for i in MaxNeta:
    NetaBest.append(MaxNeta[i])
    
NetaBest.sort()
NetaBest = NetaBest[::-1]

ans = 0
tmpans = 0
count = 0
tmpD = {}
Second = []

for i in range(K):
    d, t = ALL[i]
    if not t in tmpD:
        tmpD[t] = 1
        count += 1
    else:
        Second.append(d)
    tmpans += d

Second.sort()
#print(Second, count)
ans = tmpans + count**2
#print(NetaBest)

for i in range(count+1, min(K, NetaMax)+1):
    #print(i,i - (count+1))
    tmpans -= Second[i - (count+1)]
    tmpans += NetaBest[i-1]
    #print(tmpans + i**2)
    ans = max(ans, tmpans + i**2)
    
print(ans)

Submission status: Accepted

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []
for key in sushi:
    sushi[key].sort(reverse=True)
    first.append(sushi[key][0])
    second += sushi[key][1:]

first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))


ans = 0
for i in range(1, K+1):
    if K - i <= len(second_cum)-1 and i <= len(first_cum)-1:
        ans = max(ans, first_cum[i] + second_cum[K-i] + i*i)

print(ans)

Submission status: Accepted

Solution code:
import heapq as hq

n,k=map(int,input().split())
A=[]
S0=[] #種類が増えない寿司
S1=[] #種類が増える寿司
for i in range(n):
    t,d=map(int,input().split())
    A.append([t,d])
A=list(reversed(sorted(A)))

now=0
for i in range(n):
    t=A[i][0]
    d=A[i][1]
    if now==t:
        S0.append(-d)
    else:
        S1.append(-d)
        now=t
hq.heapify(S0)
hq.heapify(S1)
C=[] #消し得る寿司のストック(S0から出したもの)

x=0 #おいしさの合計
y=0 #種類数
for i in range(k):
    if len(S0)==0:
        x+=-hq.heappop(S1)
        y+=1
        continue
    if len(S1)==0:
        s0=-hq.heappop(S0)
        x+=s0
        hq.heappush(C,s0)
        continue
    if S1[0]<=S0[0]:
        x+=-hq.heappop(S1)
        y+=1
    else:
        s0=-hq.heappop(S0)
        x+=s0
        hq.heappush(C,s0)
ans=x+y**2

for i in range(min(len(S1),len(C))):
    c=hq.heappop(C)
    s1=-hq.heappop(S1)
    x+=s1-c
    y+=1
    ans=max(x+y**2,ans)
print(ans)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
TD=[]
for i in range(n):
  t,d=map(int,input().split())
  TD.append([t,d])
TD.sort(key=lambda x:x[1],reverse=True)
V=set()
S=[]
ans=0
for t,d in TD[:k]:
  ans+=d
  if t in V:
    S.append(d)
  else:
    V.add(t)
ans+=len(V)**2
x=ans
for t,d in TD[k:]:
  if S==[]:
    break
  if not t in V:
    x=x-S.pop()+d-len(V)**2+(len(V)+1)**2
    V.add(t)
    ans=max(x,ans)
print(ans)

Submission status: Accepted

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()
from collections import defaultdict
from itertools import accumulate

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

sushi = defaultdict(list)

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

first = []
second = []

for nums in sushi.values():
    nums.sort(reverse=True)
    for i, j in enumerate(nums):
        if i == 0:
            first.append(j)
        else:
            second.append(j)
first.sort(reverse=True)
second.sort(reverse=True)

first_cum = [0] + list(accumulate(first))
second_cum = [0] + list(accumulate(second))

ans = 0

for i in range(K, 0, -1):
    neta = min(i, len(first))
    tmp = neta * neta + first_cum[neta]

    if K - neta > len(second):
        break
    
    tmp += second_cum[K - neta]

    ans = max(ans, tmp)

print(ans)

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],reverse=True)
NEW=[1]*(n+1)
x=0
y=0
TOP=[]
SUB=[]
A=[]
for tdi in td[:k]:
    if NEW[tdi[0]]==1:
        NEW[tdi[0]]=0
        TOP.append(tdi[1])
        y+=1
    else:
        A.append(tdi[1])
    x+=tdi[1]
    
for tdi in td[k:]:
    if NEW[tdi[0]]==1:
        SUB.append(tdi[1])
        NEW[tdi[0]]=0
ans=x+y**2
for i in range(min(k-1,len(A),len(SUB))):
    x -= A.pop()
    x += SUB[i]
    y += 1
    ans = max(ans,x+y**2)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush,heappop
n,k = map(int,input().split())
kind_value = [[] for _ in range((n+1))]
for i in range(n):
    t,d = map(int,input().split())
    kind_value[t].append(d)
for li in kind_value:
    li.sort()
first_type = [(li[-1],i)for i,li in enumerate(kind_value) if li]
first_type.sort(reverse = True)
kind = 0
value_type1 = 0
second_type = []
second_type_sum = 0
second_type_cnt = 0

answer = 0
for v,i in first_type:
    kind += 1
    if kind > k:
        break
    value_type1 += v
    for d in kind_value[i][:-1]:
        heappush(second_type,d)
        second_type_sum += d
        second_type_cnt += 1
    while second_type_cnt > k-kind:
        d = heappop(second_type)
        second_type_cnt -= 1
        second_type_sum -= d
    value = kind*kind + value_type1 + second_type_sum
    if answer < value:
        answer = value

print(answer)

Submission status: Accepted

Solution code:
def main():
    import sys
    input = sys.stdin.readline
    sys.setrecursionlimit(10**7)
    from collections import Counter, deque
    #from collections import defaultdict
    from itertools import combinations, permutations, accumulate, groupby
    #from itertools import product
    from bisect import bisect_left,bisect_right
    from heapq import heapify, heappop, heappush
    from math import floor, ceil
    #from operator import itemgetter

    #inf = 10**17
    #mod = 10**9 + 7

    n,k = map(int, input().split())
    td = [list(map(int, input().split())) for _ in range(n)]
    td.sort(reverse=True, key=lambda a: a[1])
    delicious = 0
    syurui = set()
    suteru = []
    for i in range(k):
        t, d = td[i]
        delicious += d
        if  t in syurui:
            suteru.append(d)
        else:
            syurui.add(t)
    #heapify(suteru)
    suteru.sort()
    neta = len(syurui)
    res = delicious + neta**2
    for i in range(k, n):
        t, d = td[i]
        if t in syurui:
            continue
        if len(suteru)==0:
            print(res)
            exit()
        syurui.add(t)
        x = suteru.pop(0)
        delicious += d - x
        neta += 1
        res = max(res, delicious+neta**2)
    print(res)

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
def main():
    import sys
    input = sys.stdin.readline
    sys.setrecursionlimit(10**7)
    from collections import Counter, deque
    #from collections import defaultdict
    from itertools import combinations, permutations, accumulate, groupby
    #from itertools import product
    from bisect import bisect_left,bisect_right
    from heapq import heapify, heappop, heappush
    from math import floor, ceil
    #from operator import itemgetter

    #inf = 10**17
    #mod = 10**9 + 7

    n,k = map(int, input().split())
    td = [list(map(int, input().split())) for _ in range(n)]
    td.sort(reverse=True, key=lambda a: a[1])
    delicious = 0
    syurui = set()
    suteru = []
    for i in range(k):
        t, d = td[i]
        delicious += d
        if  t in syurui:
            suteru.append(d)
        else:
            syurui.add(t)
    heapify(suteru)
    neta = len(syurui)
    res = delicious + neta**2
    for i in range(k, n):
        t, d = td[i]
        if t in syurui:
            continue
        if len(suteru)==0:
            print(res)
            exit()
        syurui.add(t)
        delicious += d - heappop(suteru)
        neta += 1
        res = max(res, delicious+neta**2)
    print(res)

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
top=[0]*n
sub=[]
for _ in range(n):
    t,d=map(int,input().split())
    t-=1
    if top[t]==0:
        top[t]=d
    else:
        if d>top[t]:
            sub.append(top[t])
            top[t]=d
        else:
            sub.append(d)
top.sort(reverse=True)
top_sum=[0]
sub.sort(reverse=True)
sub_sum=[0]

for i in top:
    if i==0:
        top_sum.append(-1)
    else:
        top_sum.append(top_sum[-1]+i)
for i in sub:
    sub_sum.append(sub_sum[-1]+i)
ans=0
for i in range(1,k+1):
    if i+len(sub_sum)-1>=k and top_sum[i]!=-1:
        ans=max(ans,top_sum[i]+sub_sum[k-i]+i**2)
print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict

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 d > top[t]:
            sub.append(top[t])
            top[t] = d
        else:
            sub.append(d)

top.sort(reverse=True)
top_sum = [0]
sub.sort(reverse=True)
sub_sum = [0]

for p in top:
    if p == 0:
        top_sum.append(-1)
    else:
        top_sum.append(top_sum[-1] + p)

for p in sub:
    sub_sum.append(sub_sum[-1] + p)

ans = 0

for i in range(1, k + 1):
    if (len(sub_sum) - 1) >= (k - i) and top_sum[i] != -1:
        ans = max(ans, top_sum[i] + sub_sum[k - i] + i ** 2)

print(ans)