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:
def main():
    N, K = map(int, input().split())
    SUSHI = []
    for _ in range(N):
        t, d = map(int, input().split())
        SUSHI.append((d, t))
    SUSHI.sort(reverse=True)

    NEW = {i: True for i in range(1, N+1)} ## NEW[i]:=i番目のすしの種類がカウントされたかどうか
    x, y = 0, 0 # xは基礎ポイント、yは種類数,y*yで種類ぽいんと
    A, B = [], [] # Ａはおいしさ順に先頭K個とったときに、新種じゃないもののおいしさ
                  #　Bはおいしさ順にみたときに、Ａ以降のネタでＡに出現しなかったネタの美味しさ

    ## 以下、ＳＵＳＨＩを先頭からみて初期値x, y, A, Bを作成する
    for d, t in SUSHI[:K]:
        if NEW[t]: # 新種のネタ
            x += d
            y += 1 # 種類が増える
            NEW[t] = False
        else:
            x += d
            A.append(d)
    for d, t in SUSHI[K:]:
        if NEW[t]:
            B.append(d)
            NEW[t] = False
    B.sort()
    ### この時点で
    ### Aには先頭K種類のすしで種類ポイントに寄与しないネタ(の美味しさ)が大きい順に入っている。pop()で一番小さい値をとってこれる
    ### Bには先頭K以降のすしで種類ポイントに寄与するネタ(の美味しさ)が小さい順に入っている。pop()で一番大きい値をとってこれる。

    answer = x + y**2
    while A and B: # AかＢが空になるまで続ける
        a, b = A.pop(), B.pop() # A中種類ポイントに寄与せず、一番おいしくないネタとB中種類ポイントに寄与して、一番おいしいネタを交換
        x -= a
        x += b
        y += 1
        if x + y**2 > answer:
            answer = x + y**2
    print(answer)
main()

Submission status: Accepted

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

Submission status: Accepted

Solution code:
import heapq
n,k = map(int,input().split())
L = []
for i in range(n):
    t,d = map(int,input().split())
    L.append([d,t])
L.sort()
ans = []
cur = [0]*n
heap = []
var = 0
for i in range(k):
    if cur[L[n-i-1][1]-1] == 0:
        ans.append(L[n-i-1][0])
        cur[L[n-i-1][1]-1] = 1
        var += 1
    else:
        heapq.heappush(heap, L[n-i-1])
temp = 0
for i in range(len(ans)):
    temp += ans[i]
for i in range(len(heap)):
    temp += heap[i][0]
temp += var**2
now = temp
for i in range(n-k):
    if len(heap) != 0:
        a = heapq.heappop(heap)
        if cur[L[n-k-i-1][1]-1] == 0:
            cur[L[n-k-i-1][1]-1] = 1
            now -= var**2
            now -= a[0]
            var += 1
            now += var**2
            now += L[n-k-i-1][0]
            temp = max(temp, now)
        else:
            heapq.heappush(heap, a)
print(temp)

Submission status: Accepted

Solution code:
import heapq
n,k = map(int,input().split())
L = []
for i in range(n):
    t,d = map(int,input().split())
    L.append([d,t])
L.sort()
ans = []
cur = [0]*n
heap = []
var = 0
for i in range(k):
    if cur[L[n-i-1][1]-1] == 0:
        ans.append(L[n-i-1][0])
        cur[L[n-i-1][1]-1] = 1
        var += 1
    else:
        heapq.heappush(heap, L[n-i-1])
temp = 0
for i in range(len(ans)):
    temp += ans[i]
for i in range(len(heap)):
    temp += heap[i][0]
temp += var**2
now = temp
for i in range(n-k):
    if len(heap) != 0:
        a = heapq.heappop(heap)
        if cur[L[n-k-i-1][1]-1] == 0:
            ans.append(L[n-k-i-1][0])
            cur[L[n-k-i-1][1]-1] = 1
            now -= var**2
            now -= a[0]
            var += 1
            now += var**2
            now += L[n-k-i-1][0]
            temp = max(temp, now)
        else:
            heapq.heappush(heap, a)
print(temp)

Submission status: Accepted

Solution code:
from heapq import heappop,heappush
n,k=map(int,input().split())
s=[]
h=set()#種類
l=[0]*n #個数について
for _ in range(n):
    t,d=map(int,input().split())
    s.append((d,t-1))
s=sorted(s,reverse=True)
a=[]#heapq
num=0#美味しさのわ
for i in range(k):
    d,t=s[i]
    l[t]+=1
    h.add(t)
    num+=d
    if l[t]>1:
        heappush(a,(d,t))#２種類以上ならheapqにたす
ans=len(h)**2+num
for j in range(k,n):
    d,t=s[j]#次に小さいのを引いてくる
    l[t]+=1
    if  l[t]>=2:
        continue #2以上ならむし
    if len(a)==0:
        break
    d1,t1=heappop(a) #いらないのを引く
    h.add(t)
    num=num-d1+d
    ans=max(ans,len(h)**2+num)
print(ans)

Submission status: Accepted

Solution code:
# tokenakatta ...
def main():
    N, K = map(int, input().split())
    l = []
    for _ in range(N):
        t, d = map(int, input().split())
        l.append((t, d))
    l.sort(key=lambda x:x[1], reverse=True)
    s = set()
    e = []
    t = 0
    for i in range(K):
        t += l[i][1]
        if l[i][0] in s:
            e.append(l[i])
        else:
            s.add(l[i][0])
    rl = [t + len(s) * len(s)]
    e.reverse()
    ei = 0
    for i in range(K, N):
        if l[i][0] in s:
            continue
        if ei >= len(e):
            break
        s.add(l[i][0])
        t -= e[ei][1]
        t += l[i][1]
        ei += 1
        rl.append(t + len(s) * len(s))
    return max(rl)
print(main())

Submission status: Accepted

Solution code:
def main():
    N, K = map(int, input().split())
    TD = [tuple(map(int, input().split())) for _ in range(N)]
    TD.sort(key=lambda x: x[1], reverse=True)
    s = set()
    rep = []
    for t, d in TD[:K]:
        if t in s:
            rep.append(d)
        else:
            s.add(t)
    m = sum(d for _, d in TD[:K])
    p = len(s) ** 2 + m
    for t, d in TD[K:]:
        if not rep:
            break
        if t in s:
            continue
        m += d - rep.pop()
        s.add(t)
        p = max(p, len(s) ** 2 + m)
    return p

print(main())

Submission status: Accepted

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

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
M = len(set([t for t, d in TD]))

TD.sort(key=lambda x: x[1], reverse=True)
p, S, V = 0, [], set()
for t, d in TD[:K]:
    p += d
    if t in V:
        S.append(d)
    else:
        V.add(t)
p += len(V) ** 2
ans = p
for t, d in TD[K:]:
    if len(S) == 0:
        break
    if t not in V:
        p += 2 * len(V) + 1 - S.pop() + d
        V.add(t)
        ans = max(ans, p)
print(ans)

Submission status: Accepted

Solution code:
from collections import Counter
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)
ans = list(range(k))
ansv = sum((td[i][1] for i in range(k)))
c = Counter((td[i][0] for i in range(k)))
cnt = len(c)
ansv += len(c)**2
ptr = k
ansls = [ansv]
for i in range(k-1,-1,-1):
  t,d = td[ans[i]]
  if ptr < n and c[t] > 1:
    while ptr < n and c[td[ptr][0]] >= 1:
      ptr += 1
    if ptr >= n:
      break
    ansv = ansv-d+2*cnt+1+td[ptr][1]
    c[t] -= 1
    c[td[ptr][0]] += 1
    cnt += 1
    ptr += 1
    ansls.append(ansv)
print(max(ansls))

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
X=[]
for i in range(n):
    t,d=map(int,input().split())
    X.append([t,d])
    
X.sort(key=lambda x:x[1],reverse=True)

V=set()
point=0
S=[]

for t,d in X[:k]:
    point+=d
    if t in V:
        S.append(d)
    else:
        V.add(t)
point+=len(V)**2

ans=point

for t,d in X[k:]:
    if not S:
        break
    if not t in V:
        point+=d-S.pop()+1+len(V)*2
        V.add(t)
        ans=max(ans,point)

print(ans)

Submission status: Accepted

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

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

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

used = [0 for _ in range(N)]
dup = []
deli = 0

for i in range(K):
    t, d = sushi[i]
    if used[t - 1]:
        dup.append(d)
    used[t - 1] = 1
    deli += d

vari = sum(used)
sat = vari**2 + deli

dup.sort(reverse=True)

for i in range(N - K):
    if not dup:
        break
    t, d = sushi[K + i]
    if used[t - 1]:
        continue
    used[t - 1] = 1
    deli -= dup.pop()
    deli += d
    vari += 1
    sat = max(sat, vari**2 + deli)

print(sat)

Submission status: Accepted

Solution code:
# 写経AC

from heapq import heappop,heappush

INF = 10 ** 18

N, K = map(int, input().split())
# S[i]: 種類iの美味しさの集合
S = [[] for _ in range(N)]
for _ in range(N):
    t,d = map(int, input().split())
    S[t-1].append(d)

# 各種類について、美味しい順に並べる
for i in range(N):
    S[i].sort(reverse = True)
    if not S[i]: S[i].append(-INF) # 番兵

# 「美味しさの最大値」の降順でソート
S.sort(reverse = True)

uma = 0 # おいしさ基礎ポイント
que = [] # 選ばれていない寿司の美味しさリスト
for i in range(K):
    # ある種類の一番美味しい寿司を選ぶ
    uma += S[i][0]
    for j in range(1, len(S[i])):
        heappush(que, (-1) * S[i][j])

# 残りの寿司もリストに追加
for i in range(K, N):
    for j in range(len(S[i])):
        heappush(que, (-1) * S[i][j])

ans = uma + K * K
# 種類を減らしていく
for x in range(K - 1, 0, -1):
    v = S[x][0] # 選ばれている最小の「おいしさ」
    w = (-1) * que[0] # 選ばれていない最大の「おいしさ」
    if v < w:
        heappop(que)
        heappush(que, (-1) * v)
        uma -= v; uma += w
    ans = max(ans, uma + x * x)

print(ans)

Submission status: Accepted

Solution code:
# 解説見て書く

N, K = [int(x) for x in input().split()]
td = [0] * (N)
for i in range(N):
    td[i] = [int(x) for x in input().split()]

# ネタ種類降順でソート
# 同じネタの中ではおいしいものから順に並ぶ
td = sorted(td, reverse = True)

var_p1 = [0] * N   # 種類数を増やす寿司のリスト
var_p0 = [0] * N   # 種類数を増やさない寿司のリスト

n_p1 = 0
n_p0 = 0
for i in range(N):
    if i == 0 or (i > 0 and td[i][0] != td[i - 1][0]):  # 種類が変わる
        var_p1[n_p1] = td[i][1]  # おいしさを追加
        n_p1 += 1
    else:
        var_p0[n_p0] = td[i][1]  # おいしさを追加
        n_p0 += 1

var_p1.sort(reverse = True)
var_p0.sort(reverse = True)

val_taste = sum(var_p0[0 : K])
val_vars = 0
ans = 0

for j in range(1, K + 1):  # 種類数1〜Kを試す
    val = var_p1[j - 1]
    if val <= 0: break  # これ以上種類を増やせない
    # var_p1からn_p1個選び、var_p0からK - n_p1個選ぶ
    val_taste = val_taste + val - var_p0[K - j]
    val_vars = j * j
    ans = max(ans, val_taste + val_vars)

print(ans)

Submission status: Accepted

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


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

SUSHI = []
for _ in range(n):
    t, d = map(int, input().split())
    SUSHI.append((d, t))
SUSHI.sort(reverse = True)

NEW = {i:True for i in range(1, n + 1)}
x, y = 0, 0
A, B = [], []

for d, t in SUSHI[:k]:
    if NEW[t]:
        x += d; y += 1
        NEW[t] = False
    else:
        x += d
        A.append(d)

for d, t in SUSHI[k:]:
    if NEW[t]:
        B.append(d)
        NEW[t] = False
B.sort()

answer = x + y**2
while A and B:
    a, b = A.pop(), B.pop()
    x -= a; x += b; y += 1
    if x + y**2 > answer:
        answer = x + y**2

print(answer)

Submission status: Accepted

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

N, K = map(int, input().split())
sushi = [None] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi[i] = (d, t)
sushi.sort(reverse=True)

types = set()
cand = []
s = x = 0

for d, t in sushi[:K]:
    if t in types:
        cand.append(d)
    else:
        types.add(t)
        x += 1
    s += d
cand.sort(reverse=True)

ans = s + x*x
for d, t in sushi[K:]:
    if t in types:
        continue
    if not cand:
        break
    dr = cand.pop()
    s += d - dr
    types.add(t)
    x += 1
    ans = max(ans, s + x*x)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
import sys
input = sys.stdin.readline

N, K = map(int, input().split())
sushi = [None] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi[i] = (d, t-1)
sushi.sort(reverse=True)

types = set()
cand = []
s = x = 0

for d, t in sushi[:K]:
    if t in types:
        heappush(cand, d)
    else:
        types.add(t)
        x += 1
    s += d

ans = s + x*x
for d, t in sushi[K:]:
    if t in types:
        continue
    if not cand:
        break
    dr = heappop(cand)
    s += d - dr
    types.add(t)
    x += 1
    ans = max(ans, s + x*x)
print(ans)

Submission status: Accepted

Solution code:
import heapq
N,K = map(int,input().split())
td = [list(map(int,input().split())) for i in range(N)]
td.sort(key = lambda x:-x[1])
sumn = [[] for i in range(N)]
for i in range(N):
    heapq.heapify(sumn[i])
ans = 0
kind = {}
for t,d in td[:K]:
    if t not in kind:
        kind[t] = 1
    else:
        kind[t] += 1
    heapq.heappush(sumn[t-1],d)
kinds = len(kind)
ans = sum(sum(i) for i in sumn)
ans1 = ans +kinds**2
rem = td[K:]
nex = 0
candel = []
heapq.heapify(candel)
for k,v in kind.items():
    for i in range(v-1):
        heapq.heappush(candel,heapq.heappop(sumn[k-1]))
if candel:
    qd = heapq.heappop(candel)
    while K+nex < N:
        if td[K+nex][0] not in kind:
            kind[td[K+nex][0]] = 1
            kinds += 1
            ans = ans - qd + td[K+nex][1]
            ans1 = max(ans1,ans+kinds**2)
            nex += 1
            if candel:
                qd = heapq.heappop(candel)
            else:
                break
        else:
            nex += 1
print(ans1)

Submission status: Accepted

Solution code:
import sys

def solve():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    S = []
    for _ in range(N):
        t, d = map(int, input().split())
        S.append((d, t))
    S.sort(reverse = True)
    upperKCount = 0
    appeared = set()
    maxPoint = 0
    for i in range(K):
        maxPoint += S[i][0]
        if S[i][1] not in appeared:
            upperKCount += 1
            appeared |= {S[i][1]}
    maxPoint += upperKCount ** 2
    appeared = set()

    typeMax = []
    typeNum = 0
    typeLower = []
    others = 0
    for d, t in S:
        if t in appeared: 
            typeLower.append(d)
            others += 1
        else:
            typeMax.append(d)
            typeNum += 1
            appeared |= {t}
    typeMaxSum = [0] * (typeNum + 1)
    typeLowerSum = [0] * (others + 1)
    for i in range(typeNum): typeMaxSum[i+1] = typeMaxSum[i] + typeMax[i]
    for i in range(others): typeLowerSum[i+1] = typeLowerSum[i] + typeLower[i]
    for k in range(upperKCount + 1, min(typeNum, K) + 1):
        maxPoint = max(maxPoint, k ** 2 + typeMaxSum[k] + typeLowerSum[K - k])
    print(maxPoint)
    return 0

if __name__ == "__main__":
    solve()

Submission status: Accepted

Solution code:
N,K,*L=map(int,open(0).read().split());L=L[::-1]
I=set();X=0;U=[];D=[]
for d,t in sorted(zip(*[iter(L)]*2))[::-1]:
	if t in I:D+=d,
	elif X<K:I|={t};U+=d,;X+=1
D+=0,;i=K-X
while X and U[-1]+2*X-1<=D[i]:U.pop();X-=1;i+=1
print(sum(U+D[:i])+X**2)

Submission status: Accepted

Solution code:
N,K,*L=map(int,open(0).read().split())
S=sorted([(d,t)for t,d in zip(*[iter(L)]*2)],reverse=True)
In=set()
X=0
uni=[]
dub=[]
for d,t in S:
	if t not in In:
		In.add(t)
		uni.append(d)
		X+=1
		if X==K:
			break
	else:
		dub.append(d)
det=sum(dub[:K-X])
dub=dub[K-X:]
for d in dub:
	if uni[-1]+2*X-1<=d:
		uni.pop()
		X-=1
		det+=d
	else:
		break
print(sum(uni)+det+X**2)

Submission status: Accepted

Solution code:
N,K,*L=map(int,open(0).read().split())
S=sorted([(d,t)for t,d in zip(*[iter(L)]*2)],reverse=True)
In=set()
C=X=0
uni=[]
dub=[]
for d,t in S:
	if t not in In:
		In.add(t)
		uni.append(d)
		X+=1
		if X==K:
			break
	else:
		dub.append(d)
det=0
C=X
for d in dub:
	if C<K:
		det+=d
		C+=1
	elif uni[-1]+2*X-1<=d:
		uni.pop()
		X-=1
		det+=d
print(sum(uni)+det+X**2)

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)

types = set()
can_remove = []
d_sm = 0

for t, d in td[:k]:
	d_sm += d
	if t in types:
		can_remove.append(d)
	else:
		types.add(t)

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

left = td[k:][::-1]

while can_remove and left:
	t, d = left.pop()
	if t not in types:
		types.add(t)
		d_rmv = can_remove.pop()
		d_sm += d - d_rmv
		type_cnt += 1
		ans = max(ans, d_sm + type_cnt ** 2)

print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop

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)

types = set()
can_remove = []
d_sm = 0

for t, d in td[:k]:
	d_sm += d
	if t in types:
		can_remove.append(d)
	else:
		types.add(t)

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

left = []
for t, d in td[k:]:
	d = -d
	heappush(left, [d, t])

while can_remove and left:
	d, t = heappop(left)
	d = -d
	if t not in types:
		types.add(t)
		d_rmv = can_remove.pop()
		d_sm += d - d_rmv
		type_cnt += 1
		ans = max(ans, d_sm + type_cnt ** 2)

print(ans)

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
sushi=[]
for i in range(N):
    t,d=map(int,input().split())
    sushi.append((t,d))
sushi.sort(key=lambda x:(x[0],-x[1]))
#print(sushi)
X=[sushi[0][1]]
Y=[]
for i in range(1,N):
    if sushi[i-1][0]==sushi[i][0]:
        Y.append(sushi[i][1])
    else:
        X.append(sushi[i][1])
X.sort(reverse=True)
Y.sort(reverse=True)
SX=[0]
SY=[0]
for i in X:
    SX.append(SX[-1]+i)
for i in Y:
    SY.append(SY[-1]+i)
ans=0
#print(X)
#print(Y)
#1<=x<len(X)
#K-len(Y)<=x<=
for x in range(max(1,K-len(Y)),min(len(SX),K+1)):
    ans=max(SX[x]+SY[K-x]+x*x,ans)
print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.buffer.readline
import heapq
from collections import defaultdict
from operator import itemgetter

def main():
    N,K = map(int,input().split())
    sushi = []
    for _ in range(N):
        t,d = map(int,input().split())
        sushi.append((-d,t))
    sushi.sort(key=itemgetter(0))
    use,only,unuse = [],[],sushi[K:]
    neta = defaultdict(int)
    x = 0
    for i in range(K):
        d,t = sushi[i]
        if neta[t] == 0:
            only.append((-d,t))
        else:
            use.append((-d,t))
        neta[t] += 1
        x -= d
    
    l = len(neta)
    heapq.heapify(use)
    heapq.heapify(only)
    heapq.heapify(unuse)
    
    ans = x + l**2
    while (len(use) != 0 and len(unuse) != 0 and l < K):
        cd,ct = heapq.heappop(unuse)
        cd *= -1
        if neta[ct] != 0:
            continue
        heapq.heappush(only,(cd,ct))
        rd,rt = heapq.heappop(use)
        x += cd-rd
        l += 1
        neta[ct] += 1
        neta[rt] -= 1
        ans = max(ans,x+l**2)
        
    print(ans)

if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
N, K, *L = map(int, open(0).read().split())
ls = []
for t,d in zip(*[iter(L)]*2):
  ls += [(d,t)]
used = []
cnt = [0]*(N+1)
hq = []
S = set()
ls.sort(reverse=True)
m = 0
for i in range(K):
  d,t = ls[i]
  cnt[t] += 1
  m += d
  if cnt[t]>1:
    used.append(d)
  S.add(t)
used.sort(reverse=True)
m += len(S)**2
for i in range(K,N):
  d,t = ls[i]
  if t not in S:
    heappush(hq,(-d,t))
ans = [m]
while used:
  d = used.pop()
  while hq:
    if hq[0][1] not in S:
      break
    heappop(hq)
  else:
    break
  m -= d
  n,s = heappop(hq)
  m += -n+2*len(S)+1
  ans.append(m)
  S.add(s)
print(max(ans))

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
N, K, *L = map(int, open(0).read().split())
ls = []
for t,d in zip(*[iter(L)]*2):
  ls += [(d,t)]
used = []
cnt = {}
hq = []
S = set()
ls.sort(reverse=True)
m = 0
for i in range(K):
  d,t = ls[i]
  cnt[t] = cnt.get(t,0)+1
  m += d
  if cnt[t]>1:
    heappush(used,(d,t))
  S.add(t)
m += len(S)**2
for i in range(K,N):
  d,t = ls[i]
  if t not in S:
    heappush(hq,(-d,t))
ans = m
while used:
  d,t = heappop(used)
  while hq:
    if hq[0][1] not in S:
      break
    heappop(hq)
  else:
    break
  m -= d
  n,s = heappop(hq)
  m += -n+2*len(S)+1
  ans = max(ans,m)
  S.add(s)
print(ans)

Submission status: Accepted

Solution code:
import sys
import heapq as h

input=sys.stdin.readline

N,K=map(int,input().split())
sushi=[[0,0] for i in range(0,N)]
for i in range(0,N):
    t,k=map(int,input().split())
    sushi[i][1]=t
    sushi[i][0]=k

sushikind=[[] for i in range(0,N+1)]
for i in range(0,N):
    sushikind[sushi[i][1]].append(sushi[i][0])

kindmax=[0 for i in range(0,N+1)]

for i in range(0,N):
    if sushikind[i]:
        kindmax[i]=max(sushikind[i])

sushi.sort(reverse=True)
sum=0
kind=[]
kind=set(kind)
extra=[]
flag=[0 for i in range(0,N+1)]
for i in range(0,K):
    sum+=sushi[i][0]
    kind.add(sushi[i][1])
    if sushi[i][0]==kindmax[sushi[i][1]]:
        if flag[sushi[i][1]]==1:
            extra.append(sushi[i][0])
        else:
            flag[sushi[i][1]]=1
    else:
        extra.append(sushi[i][0])


set={i for i in range(1,N+1)}
set-=kind
appendix=[kindmax[i] for i in set if kindmax[i]!=0 ]
extra.sort()
appendix.sort(reverse=True)
k=len(kind)
first=k
ans=k**2+sum
for i in range(k,min(K,k+len(appendix))):
    k+=1
    sum+=appendix[k-first-1]-extra[k-first-1]
    test=k**2+sum
    if test>ans:
        ans=test

print(ans)

Submission status: Accepted

Solution code:

n,k= map(int, input().split())
t= [list(map(int, input().split())) for i in range(n)]
q=[]
ans=0
r=[0]*(n+1)

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

for i in range(k):
    ans += t[i][1]
    if r[t[i][0]]==0:
        r[t[i][0]]=1
    else:
        q.append(t[i][1])
v=sum(r)
ans+=v**2
cnt=ans


for i in range(k,n):
    if r[t[i][0]]==0 and len(q)!=0:
        r[t[i][0]]=1
        x=q.pop()
        cnt+=t[i][1]+2*v+1-x
        v+=1
        ans=max(cnt,ans)
print(ans)

Submission status: Accepted

Solution code:
import collections
from collections import deque
n,k= map(int, input().split())
t= [list(map(int, input().split())) for i in range(n)]
q=deque()
ans=0
r=[0]*(n+1)

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

for i in range(k):
    ans += t[i][1]
    if r[t[i][0]]==0:
        r[t[i][0]]=1
    else:
        q.append(t[i][1])
v=sum(r)
ans+=v**2
cnt=ans


for i in range(k,n):
    if r[t[i][0]]==0 and len(q)!=0:
        r[t[i][0]]=1
        x=q.pop()
        cnt+=t[i][1]+2*v+1-x
        v+=1
        ans=max(cnt,ans)
print(ans)

Submission status: Accepted

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

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q1 = dd(list)
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

ans = score + kind*kind
while q2 and q3:
    score -= heappop(q2)+heappop(q3)
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Submission status: Accepted

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

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q1 = dd(list)
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

ans = score + kind*kind
while q2 and q3:
    d1 = heappop(q2)
    d2 = heappop(q3)
    score -= d1+d2
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Submission status: Accepted

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

N, K = map(int, input().split())
l = [tuple(map(int, input().split())) for _ in range(N)]
l.sort(key=lambda x:-x[1])
used = set()
q2, q3 = [], []

score, kind = 0, 0

for t, d in l[:K]:
    score += d
    if t in used:
        heappush(q3, d)
    else:
        used.add(t)
        kind += 1

ans = score + kind*kind

q1 = dd(list)
for t, d in l[K:]:
    if not t in used:
        heappush(q1[t], -d)
for t in q1.keys():
    d = heappop(q1[t])
    heappush(q2, d)

while q2 and q3:
    d1 = heappop(q2)
    d1 = -d1

    d2 = heappop(q3)
    score -= d2
    score += d1
    kind += 1
    ans = max(ans, score + kind*kind)

print(ans)

Submission status: Accepted

Solution code:
n,m = map(int,input().split())
a = []
for i in range(n):
    r,t = map(int,input().split())
    a.append((t,r))
a.sort(reverse = True)
#print(a)
d = {}
cd = []
dp = 0
x = 0
 
for i in range(m):
    z,c = a[i]
    dp += z
    if c in d:
        cd.append(z)
    else:
        d[c] = z
        x += 1
#print(cd)
cd.sort(reverse = True)
ans = dp + x**2
c = ans
 
for i in range(m,n):
    if len(cd) == 0:
        break
    z,v = a[i]
    if v in d:
        continue
    d[v] = z
    x += 1
    c = c - cd.pop()+z+2*x-1
    if c > ans:
        ans = c
print(ans)

Submission status: Accepted

Solution code:
import sys
def input():
    return sys.stdin.readline()[:-1]
N,K=map(int, input().split())
s=[None]*N
for i in range(N):
    t,d=map(int,input().split())
    s[i]=(d,t)
s.sort(reverse=True)
l=[0]*N
d=0
for i in s[:K]:
    l[i[1]-1]+=1
    d+=i[0]
c=0
for i in l:
    if i:
        c+=1
a=d+c**2
h=K-1
for i in s[K:]:
    if not l[i[1]-1]:
        while h>=0:
            if l[s[h][1]-1]>1:
                l[i[1]-1]=1
                l[s[h][1]-1]-=1
                c+=1
                d=d+i[0]-s[h][0]
                a=max(a,d+c**2)
                h-=1
                break
            h-=1
        else:
            break
print(a)

Submission status: Accepted

Solution code:
import sys

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

N, K = lr()
TD = [lr() for _ in range(N)]
TD.sort(key=lambda x: x[1], reverse=True)
used = set()
overlap = []
for t, d in TD[:K]:
    if t in used:
        overlap.append(d)
    else:
        used.add(t)

delicious = sum([x[1] for x in TD[:K]])
answer = delicious + len(used) ** 2
for i in range(K, N):
    if not overlap:
        break
    if TD[i][0] in used:
        continue
    delicious += (TD[i][1] -  overlap.pop())
    used.add(TD[i][0])
    if answer < delicious + len(used)**2:
        answer = delicious + len(used)**2

print(answer)
# 30

Submission status: Accepted

Solution code:
import heapq

N, K = map(int, input().split())
dt = [list(map(int, input().split()))[::-1] for i in range(N)]
dt.sort(reverse=True)

ans = 0
t_set = set()
candidates = []
for d, t in dt[:K]:
  if t in t_set:
    heapq.heappush(candidates, [d, t])
    ans += d
  else:
    t_set.add(t)
    ans += d

x = len(t_set)
ans += x**2

if len(candidates) == 0:
  print(ans)
  quit()

ans = [ans]
for d, t in dt[K:]:
  if t in t_set:
    continue
  cand_d, cand_t = heapq.heappop(candidates)
  _ans = ans[-1] - cand_d + d + (2*x+1)
  x += 1
  ans.append(_ans)
  t_set.add(t)
  if len(candidates) == 0:
    break

print(max(ans))

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
l=[]
for i in range(n):
    l.append(list(map(int,input().split())))
l.sort()
l.append([-1,-1])
a=[]
b=[]
for i in range(n):
    if l[i][0]==l[i+1][0]:
        b.append(l[i][1])
    else:
        a.append(l[i][1])
ans=len(a)*len(a)+sum(a)
#print(ans)
if len(b)>0:
    ans+=sum(b)
    b.sort()
    a.sort()
    i=0
    ai=0
    bi=0
    an=len(a)
    bn=len(b)
    while i<n-k:
        if ai<an and bi<bn:
            ap=a[ai]
            bp=b[bi]
            if ap+2*(an-ai-1)+1>=bp:
                bi+=1
                ans-=bp
            else:
                ans-=ap+2*(an-ai-1)+1
                ai+=1
        elif ai<an:
            ap=a[ai]
            ans-=ap+2*(an-ai-1)+1
            ai+=1
        elif bi<bn:
            ans-=b[bi]
            bi+=1
        i+=1
    print(ans)
else:
    a.sort(reverse=True)
    print(sum(a[:k])+k*k)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
l=[]
for i in range(n):
    l.append(list(map(int,input().split())))
l.sort(key=lambda x: x[0]*1000000000000+x[1])
l.append([-1,-1])
a=[]
b=[]
for i in range(n):
    if l[i][0]==l[i+1][0]:
        b.append(l[i][1])
    else:
        a.append(l[i][1])
ans=len(a)*len(a)+sum(a)
#print(ans)
if len(b)>0:
    ans+=sum(b)
    b.sort()
    a.sort()
    i=0
    ai=0
    bi=0
    an=len(a)
    bn=len(b)
    while i<n-k:
        if ai<an and bi<bn:
            ap=a[ai]
            bp=b[bi]
            if ap+2*(an-ai-1)+1>=bp:
                bi+=1
                ans-=bp
            else:
                ans-=ap+2*(an-ai-1)+1
                ai+=1
        elif ai<an:
            ap=a[ai]
            ans-=ap+2*(an-ai-1)+1
            ai+=1
        elif bi<bn:
            ans-=b[bi]
            bi+=1
        i+=1
    print(ans)
else:
    a.sort(reverse=True)
    print(sum(a[:k])+k*k)

Submission status: Accepted

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

Submission status: Accepted

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

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

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

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

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

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

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

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

print(ans)