Problem:


Score : 400 points

Problem Statement

There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i.
You are choosing K among these N pieces to eat.
Your "satisfaction" here will be calculated as follows:

The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".

The base total deliciousness is the sum of the deliciousness of the pieces you eat.

The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.

You want to have as much satisfaction as possible.
Find this maximum satisfaction.

Constraints

1 \leq K \leq N  \leq 10^5

1 \leq t_i \leq N

1 \leq d_i \leq 10^9

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

Output

Print the maximum satisfaction that you can obtain.

Sample Input 1

5 3
1 9
1 7
2 6
2 5
3 1

Sample Output 1

26

If you eat Sushi 1,2 and 3:

The base total deliciousness is 9+7+6=22.

The variety bonus is 2*2=4.

Thus, your satisfaction will be 26, which is optimal.

Sample Input 2

7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5

Sample Output 2

25

It is optimal to eat Sushi 1,2,3 and 4.

Sample Input 3

6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000

Sample Output 3

4900000016

Note that the output may not fit into a 32-bit integer type.



Sample input:
5 3
1 9
1 7
2 6
2 5
3 1


Sample output:
26

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
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)
A=td[:K]
B=[]
s2=set()
for i,j in td[K:]:
    if not(i in s2):
        s2.add(i)
        B.append([i,j])
s1=set()
n=0
C=[]
sum=0
for i,j in A:
    if not(i in s1):
        s1.add(i)
        n+=1
        sum+=j
    else:
        C.append([i,j])
        sum+=j
sum+=n**2
B.reverse()
D=[]
for i,j in B:
    if not(i in s1):
        D.append([i,j])
list=[sum]

while(C):
    if D:
        i,j=D.pop()
        list.append(list[-1]+2*n+1+(j-C[-1][1]))
        n+=1
    C.pop()
print(max(list))

Submission status: Accepted

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

max_kiso = sum([d for t,d in TD[:K]])
classes = set([t for t,d in TD[:K]])
n_classes = len(classes)

from itertools import groupby, accumulate
to_remove = []
for k,g in groupby(sorted(TD[:K]), key = lambda x:x[0]):
    #k: t
    to_remove.extend(list(g)[:-1]) #最大のもの以外
d_to_remove = sorted([d for t,d in to_remove])
d_to_remove_cumsum = list(accumulate(d_to_remove))

assert n_classes + len(to_remove) == K

to_add = []
for k,g in groupby(sorted(TD[K:]), key = lambda x:x[0]):
    if k not in classes:
        to_add.append(list(g)[-1]) #最大のもの
d_to_add = sorted([d for t,d in to_add], reverse= True)
d_to_add_cumsum = list(accumulate(d_to_add))

answer = max_kiso + n_classes **2
for i in range(0, min(len(to_remove), len(to_add))):
    a = max_kiso + d_to_add_cumsum[i] - d_to_remove_cumsum[i] + (n_classes + i+1)**2
    answer = max(a,answer)

print(answer)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])
ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans1 + TD[k1][1] - q.pop() - len(s) ** 2 + (len(s) + 1) ** 2
        ans = max(ans, ans1)
        s.add(TD[k1][0])
    k1 += 1
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)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])
ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans1 + TD[k1][1] - q.pop() - len(s) ** 2 + (len(s) + 1) ** 2
        ans = max(ans, ans1)
        s.add(TD[k1][0])
    k1 += 1
print(ans)

Submission status: Accepted

Solution code:
n,k=map(int, input().split())
a=sorted([list(map(int, input().split())) for i in range(n)], key=lambda x:-x[1])
q=[]
ans=0
s=set()
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
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 heapq import*
I=lambda:list(map(int,input().split()))[::-1]
k,n=I()
q=[]
v=i=T=A=0
l=[0]*-~n
for b,a in sorted(eval('I(),'*n))[::-1]:
 if i<k:
  if l[a]:heappush(q,b)
  else:T+=v*2+1;v+=1
  T+=b
 elif l[a]^1and q:T+=b+v*2+1-heappop(q);v+=1
 A=max(A,T);i+=1;l[a]=1
print(A)

Submission status: Accepted

Solution code:
import sys
import heapq
#入力
n,k=map(int,input().split())
sushi=[[int(i) for i in l.split()] for l in sys.stdin]
#まずsushiを二番目の要素で降順にsort
sushi.sort(key=lambda x: x[1],reverse=True)
#次に大きい順からk個取ってきてsummationに加えると共に種類数を記録する
#この時にどの種類の寿司が加えられたかを記録しておき、二個目以降ではheapqにつっこむ
summation=0
various=0
stack=[0]*(n+1)
heap_min=[]
for i in range(k):
  if stack[sushi[i][0]]==0:
    stack[sushi[i][0]]=1
    various+=1
  else:
    heap_min.append(sushi[i][1])
  summation+=sushi[i][1]
#残ったものの中からまだ出てきていない種類の寿司の最大値をheapqにつっこむ
#このときすでにソートされているので、前からやっていってはじめに出てきたやつが最大値
heap_max=[]
for i in range(k,n):
  if stack[sushi[i][0]]==0:
    stack[sushi[i][0]]=1
    heap_max.append((-sushi[i][1],sushi[i][1]))
#heapify
heapq.heapify(heap_min)
heapq.heapify(heap_max)
#ansを計算していく
ans=summation+various*various
for i in range(1,k-various+1):
  if len(heap_min)==0 or len(heap_max)==0:
    break
  summation-=heapq.heappop(heap_min)
  summation+=heapq.heappop(heap_max)[1]
  ans=max(ans,summation+(various+i)*(various+i))
#出力
print(ans)

Submission status: Accepted

Solution code:
N,K=map(int,input().split())
Sushi=[[int(i) for i in input().split()] for j in range(N)]
Sushi=sorted(Sushi,key=lambda x:x[1],reverse=True)
inp={}
for t,d in Sushi:
    if not t in inp:
        inp[t]=d#それぞれのネタの最大値が入る
tasty=0
neta=set()#何種類か数える
dup=[]
for i in range(K):
    t,d=Sushi[i]
    tasty+=d
    if t in neta:
        dup.append((t,d))#値が大きい順に入る
    neta.add(t)
    if t in inp:
        del inp[t]
variety=len(neta)
max_tasty=tasty+variety**2
for s,t in sorted(inp.items(),key=lambda x:x[1],reverse=True):
    if len(dup):
        s1,t1=dup.pop()#一番値が小さい要素が出てくる
        variety+=1
        tasty=tasty-t1+t
        max_tasty=max(max_tasty,tasty+variety**2)
print(max_tasty)

Submission status: Accepted

Solution code:
import heapq

INF = float('inf')
n, k = map(int, input().split())
sushi = {}

for _ in range(n):
    t, d = map(int, input().split())
    if t not in sushi:
        sushi[t] = []
    sushi[t].append(d)

for s in sushi.values():
    s.sort(reverse=True)

eat_counts = {s: 0 for s in sushi.keys()}

# まだ食べてないやつ
news = []
# リピート
repeats = []
# 食べた履歴
eats = []

# 各寿司の一番おいしいやつを heap へ
for name, deli_li in sushi.items():
    # (-おいしさ, name)
    heapq.heappush(news, (-deli_li[0], name))

sum_deli = 0
variety = 0

for _ in range(k):
    # おいしいやつから食べる
    if news and (not repeats or news[0] <= repeats[0]):
        deli, name = heapq.heappop(news)
        variety += 1
    else:
        deli, name = heapq.heappop(repeats)

    sum_deli += -deli
    eat_counts[name] += 1
    heapq.heappush(eats, (-deli, name))
    if eat_counts[name] < len(sushi[name]):
        heapq.heappush(repeats, (-sushi[name][eat_counts[name]], name))

# 美味しいやつから順に食べた
# 種類増やしたらもっと得点上がるか確かめる
ans = sum_deli + variety ** 2
while True:
    if not news or not eats:
        break

    # 食べたやつポイント小さいやつから置き換えたい
    ate_min_deli, ate_min_name = heapq.heappop(eats)
    eat_counts[ate_min_name] -= 1
    if eat_counts[ate_min_name] <= 0:
        # これは種類減っちゃうからだめ
        continue

    sum_deli -= ate_min_deli

    deli, name = heapq.heappop(news)
    variety += 1
    sum_deli += -deli
    eat_counts[name] += 1

    ans = max(ans, sum_deli + variety ** 2)

print(ans)

Submission status: Accepted

Solution code:
num_lines, max_num = map(int, input().split())

inp_sorted = [[int(st) for st in input().split()] for _ in range(num_lines)]
# print(inp_sorted)
inp_sorted.sort(key=lambda x: x[1], reverse=True)

inp_map = {}
for t, d in inp_sorted:
    if not t in inp_map:
        inp_map[t] = d

variety = set()
score = 0

dup = []
for i in range(max_num):
    t, d = inp_sorted[i]
    if t in variety:
        dup.append((t,d))
    variety.add(t)
    if t in inp_map:
        del inp_map[t]
    score += d

var_value = len(variety)
max_score = score + var_value*var_value

for t, d in sorted(inp_map.items(), key=lambda x: x[1], reverse=True):
    if (len(dup)):
        t2, d2 = dup.pop()
        var_value += 1
        score = score - d2 + d
        max_score = max(max_score, score + var_value*var_value)

print(max_score)

Submission status: Accepted

Solution code:
import heapq

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

for _ in range(N):
    (t,d)=map(int,input().split())
    heapq.heappush(L,(-d,t))

LL=[]
sum=0
used=set()

for i in range(K):
    d,t=heapq.heappop(L)
    if t in used:
        heapq.heappush(LL,(-d,t))
    used.add(t)
    sum+=-d

ans=len(used)**2+sum

while L and LL:
    d,t=heapq.heappop(L)
    if t not in used:
        used.add(t)
        dd,_=heapq.heappop(LL)
        sum=sum-dd+(-d)
        ans=max(ans,len(used)**2+sum)

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]
a = [] #種類があぶれたやつを保管するリスト
s = set()
ans,cur = 0,0
for i,(d,t) in enumerate(dt,1):
    if t in s:
        if i <= K:
            a.append(d)
            cur += d
    else:
        if i > K:
            if not a:
                break
            cur -= a.pop() #種類があぶれていて価値の最も低い寿司を抜く
        cur += d
        s.add(t)
    ans = max(ans,len(s)**2+cur) #f(i-1)とf(i)の値を比べていく
print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict
import heapq

def main():
	n, k = map(int, input().split())
	td = [0]*n
	group = defaultdict(list)
	for i in range(n):
		td[i] = list(map(int, input().split()))
		group[td[i][0]].append(td[i][1])

	grouplist = list(group.values())
	for g in grouplist:
		g.sort(reverse=True)
	grouplist.sort(key=lambda x: -x[0])

	non_maxes = []
	result = 0
	base = 0

	# i stands for total num of netas - 1
	for i in range(min(len(grouplist), k)):
		bonus = (i+1)*(i+1)
		# only use grouplist[0] to grouplist[i]
		for j, basic in enumerate(grouplist[i]):
			# new group we should add at least one neta
			if j == 0:
				if len(non_maxes) > k - (i+1):
					remove = heapq.heappop(non_maxes)
				else:
					remove = 0
				# always add each max value
				base += basic - remove
			elif len(non_maxes) < k - (i+1):
				heapq.heappush(non_maxes, basic)
				base += basic
			# new groups non-max value is larger than prior groups non-max value
			elif non_maxes and non_maxes[0] < basic:
				base += basic - non_maxes[0]
				heapq.heapreplace(non_maxes, basic)
		result = max(result, base+bonus)
	print(result)


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
#!usr/bin/env python3
from collections import defaultdict
from collections import deque
from heapq import heapify,heappush, heappop
import sys
import math
import bisect
import random
def LI(): return list(map(int, sys.stdin.readline().split()))
def I(): return int(sys.stdin.readline())
def LS():return list(map(list, sys.stdin.readline().split()))
def S(): return list(sys.stdin.readline())[:-1]
def IR(n):
    l = [None for i in range(n)]
    for i in range(n):l[i] = I()
    return l
def LIR(n):
    l = [None for i in range(n)]
    for i in range(n):l[i] = LI()
    return l
def SR(n):
    l = [None for i in range(n)]
    for i in range(n):l[i] = S()
    return l
def LSR(n):
    l = [None for i in range(n)]
    for i in range(n):l[i] = LS()
    return l
sys.setrecursionlimit(1000000)
mod = 1000000007

#A
def A():
    return

#B
def B():
    return

#C
def C():
    return

#D
def D():
    n,k = LI()
    g = LIR(n)
    g.sort(key = lambda x:-x[1])
    f = [0 for i in range(n+1)]
    x = 0
    ans = 0
    c = []
    for i in range(k):
        ans += g[i][1]
        if not f[g[i][0]]:x+=1
        else:c.append(g[i][1])
        f[g[i][0]] += 1
    ans += x**2
    c.sort()
    c = deque(c)
    d = ans
    for i in range(k,n):
        if not c:
            break
        if f[g[i][0]]:continue
        d += g[i][1]-c.popleft()+2*x+1
        f[g[i][0]] = 1
        x += 1
        if d > ans:
            ans = d
    print(ans)
    return

#E
def E():
    return

#F
def F():
    return

#G
def G():
    return

#H
def H():
    return

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

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
import sys, heapq


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    group = {}
    for i in range(N):
        k = t[i]
        v = d[i]
        if k in group:
            group[k].append(v)
        else:
            group[k] = [v]
    groups = list(group.values())
    for g in groups:
        g.sort(reverse=True)
    groups.sort(reverse=True, key=lambda x: x[0])
    h = []
    ret = 0
    base = 0
    #print(groups)
    for i in range(min(len(groups), K)):
        g = groups[i]
        bonus = (i + 1) ** 2
        for j, v in enumerate(g):
            if j == 0:
                if len(h) > K - (i + 1):
                    remove = heapq.heappop(h)
                else:
                    remove = 0
                base = base - remove + v
            elif len(h) < K - (i + 1):
                heapq.heappush(h, v)
                base += v
            elif len(h) > 0 and h[0] < v:
                base += v - h[0]
                heapq.heapreplace(h, v)
            ret = max(ret, base + bonus)
            #print(h)
            #print("bonus : %d, base : %d, ret : %d" % (bonus, base, ret))
            #print(' --- ')
    print(ret)
    return


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:
sn, tn, *s = map(int, open(0).read().split())
sy_oi = [[s[i * 2], s[i * 2 + 1]] for i in range(sn)]
sy_oi.sort(key=lambda x: -x[1])
cho = sy_oi[tn - 1::-1]
zan = sy_oi[tn:]
zn = sn - tn
# チョイスした種類の辞書
sydic = {}
for x in cho:
    if x[0] in sydic:
        sydic[x[0]] += 1
    else:
        sydic[x[0]] = 1
ot = sum([x[1] for x in cho])  # おいしさポイント
mt = len(sydic)  # まんぞくポイント
ans = ot + mt ** 2
ci = 0
zi = 0
while 1:
    # チョイスから削除できるものを探す
    for i in range(ci, tn):
        if sydic[cho[i][0]] > 1:
            sydic[cho[i][0]] -= 1
            ot -= cho[i][1]
            ci = i + 1
            break
    else:
        break
    # 残から加えるものを探す
    for i in range(zi, zn):
        if zan[i][0] not in sydic:
            sydic[zan[i][0]] = 1
            ot += zan[i][1]
            mt += 1
            zi = i + 1
            break
    else:
        break
    # 得点計算して更新
    ans = max(ans, ot + mt ** 2)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
n, k =map(int, input().split())

#重み付きのヒープに入れる
L =[]
for i in range(n):
    t, d = map(int, input().split())
    heappush(L, (-d, t)) #降順にしたいけど,heapqが昇順固定なので、-にする

LL=[]
sum = 0
s = set()
for i in range(k):
    d, t = heappop(L)
    if t in s:
        heappush(LL, (-d, t)) #重複したものは交換対象としておく
    s.add(t)
    sum+=-d

ans = len(s)**2+sum

while L and LL:
    d, t = heappop(L)
    if t not in s:
        s.add(t)
        dd, _ = heappop(LL)
        sum = sum -dd +(-d)
        ans = max(ans, len(s)**2+sum)
print(ans)

Submission status: Accepted

Solution code:
from collections import Counter
N, K =  map(int, (input().split(' ')))
sushi = [tuple(map(int, input().split(' '))) for _ in range(N)]
all_sushi = sorted(sushi, key=lambda x: -x[1])
default_sushi = all_sushi[:K]
amari = all_sushi[K:]
len_amari = len(amari)
counter = Counter([i[0] for i in default_sushi])
result = []
def calc_value(sushi):
    return sum([i[1] for i in sushi]) + len(counter)**2

fs = []
def change_sushi(sushi):
    oishisa = sum([i[1] for i in sushi])
    fs.append(oishisa + len(counter)**2)
    amari_index = 0
    for s in sushi[::-1]:
        if counter[s[0]] >= 2:
            while amari_index < len_amari:
                a = amari[amari_index]
                amari_index += 1
                if a[0] not in counter:
                    counter[s[0]] -= 1
                    counter[a[0]] += 1
                    oishisa = oishisa - s[1] + a[1]
                    fs.append(oishisa + len(counter)**2)
                    break
    print(max(fs))
    exit()
            
    
def get_sushi(sushi):
    #result.append(calc_value(sushi))
    changed_sushi = change_sushi(sushi)
    #if not changed_sushi:
        #print(max(result))
        #return
    #get_sushi(changed_sushi)

get_sushi(default_sushi)

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)
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
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:
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)
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
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:
N, K = map(int, input().split())
best = [0] * N
not_best = []
for _ in range(N):
    t, d = map(int, input().split())
    if best[t-1] != 0:
        not_best.append(min(d, best[t-1]))
    best[t-1] = max(d, best[t-1])
best = [x for x in best if x != 0]
best.sort(reverse=True)
not_best.sort(reverse=True)
acc_best = [0] * (N + 1)
acc_not_best = [0] * (N + 1)
for i in range(len(best)):
    acc_best[i+1] = acc_best[i] + best[i]
for i in range(len(not_best)):
    acc_not_best[i+1] = acc_not_best[i] + not_best[i]
cand = []
for x in range(1, K+1):
    if x > len(best):
        break
    if K - x > len(not_best):
        continue
    cand.append(acc_best[x] + acc_not_best[K-x] + x**2)
print(max(cand))

Submission status: Accepted

Solution code:
import collections

index = 0

def maxsushi(sushi_list,eaten_sushi):
    global index
    max_val = [0,0]
    for i in range(index,len(sushi_list)):
        if(sushi_list[i][1] not in eaten_sushi):
            tmp = sushi_list[i]
            max_val = sushi_list[i]
            index = i
            break
    return max_val

N,K = map(int,input().split())
sushi = []
umasa = []
for _ in range(N):
    i,j = map(int,input().split())
    sushi.append(i)
    umasa.append(j)
l = list(zip(umasa,sushi))
l = sorted(l,reverse=True)
sushi_list = list(l)
eat_sushi = sushi_list[0:K]
eaten_sushi = set()
double_sushi = []
for i in range(K):
    if(sushi_list[i][1] not in eaten_sushi):
        eaten_sushi.add(sushi_list[i][1])
    else:
        double_sushi.append((sushi_list[i]))
manzoku = 0
sushi_list = sushi_list[K:N]
for i in range(len(eat_sushi)):
    manzoku += eat_sushi[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku 
while double_sushi:
    tmp1 = double_sushi.pop()
    tmp2 = maxsushi(sushi_list,eaten_sushi)
    if(tmp2[0] == 0 and tmp2[1] == 0):
        break
    eaten_sushi.add(tmp2[1])
    tmp_manzoku = tmp_manzoku + tmp2[0] - tmp1[0] + len(eaten_sushi) ** 2 - (len(eaten_sushi) - 1) ** 2
    if (tmp_manzoku > manzoku):
        manzoku = tmp_manzoku
print(manzoku)

Submission status: Accepted

Solution code:
# /usr/bin/python
# -*- coding: utf-8 -*-
import sys


n,k = map(int, input().split())
td_mx = {}
td = []
for _ in range(n):
  t,d = map(int, input().split())
  if t not in td_mx:
    td_mx[t] = d
  elif td_mx[t] < d:
    td.append(td_mx[t])
    td_mx[t] = d
  else:
    td.append(d)
td = sorted(td, key=lambda x:-x)
td_mx = sorted(td_mx.values(), key=lambda x:-x)

# td,td_mx の累積和を求める
csum = [0]
csum_mx = [0]
for i in range(len(td)):
  csum.append(csum[-1] + td[i])
for i in range(len(td_mx)):
  csum_mx.append(csum_mx[-1] + td_mx[i])


# 1~k種類で探索 最大値を取り出す
ans = 0
for i in range(1,k+1):
  if i <= len(csum_mx)-1 and k-i <= len(csum)-1:
    ans = max(ans, csum_mx[i] + csum[k-i] + i*i)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
n, k = map(int, input().split())

H = []
for i in range(n):
    t, d = map(int, input().split())
    heappush(H, (- d, t))

Q = []
used = {}
dsum = 0
for i in range(k):
    d, t = heappop(H)
    if t not in used:
        used[t] = 1
    else:
        heappush(Q, (- d, t))
    dsum += - d
    
ans = len(used) ** 2 + dsum
        
#print(H)
#print(Q)
#print(used)
#print(dsum)
#print(ans)

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

Submission status: Accepted

Solution code:
from collections import defaultdict


def main():
    """
    N個の寿司があり、それぞれの寿司には
    「ネタ」ti
    「おいしさ」di
    のパラメータが設定されています。 
    あなたはこの N個の寿司の中から K個を選んで食べようとしています。
    この時のあなたの「満足ポイント」は、以下のようにして計算されます。

    「満足ポイント」は、「おいしさ基礎ポイント」と、「種類ボーナスポイント」の和である。
    「おいしさ基礎ポイント」は、食べた寿司の「おいしさ」の総和である。
    「種類ボーナスポイント」は、食べた寿司の「ネタ」の種類数を xとしたとき、x ∗ x である。

    あなたは、「満足ポイント」をできるだけ大きくしたいです。
    この時の「満足ポイント」の値を求めてください。

    1 ≦ K  ≦ N ≦ 10^5
    1 ≦ ti ≦ N
    1 ≦ di ≦10^9

    nCk の組合せ: 10^5 C 10^5/2 => TLE
    DP ?
        1-Nまで, 食べる食べないによるスコア
        食べたとして種類数xとして種類は別途持っておく？
        →二次元配列は無理っぽそう, 10^5 * 10^5 = 10^10
        MAX時の種類だけ持って、１次元にする
        K までになっているかの判定
        →Nとは別に持たないといけない
    """
    N, K = map(int, input().split())
    td = [
        list(map(int, input().split()))
        for _ in range(N)
    ]
    ans = editorial(N, K, td)
    print(ans)


def editorial(N, K, td):
    td = sorted(td, key=lambda x: -x[1])
    c = defaultdict(int)
    d_sum = 0
    # 種類数が増えないので、不要と判断して良い基礎ポイント
    td_not_first = []
    for t, d in td[:K]:
        if c[t] > 0:
            td_not_first.append(d)
        c[t] += 1
        d_sum += d

    t_sum = len(c)
    ans = d_sum + t_sum ** 2

    for t, d in td[K:]:
        # 更新の余地なし
        if not td_not_first:
            return ans
        # すでに存在するなら、基礎ポイント降順のため不要
        if c[t] > 0:
            continue

        c[t] += 1
        d_sum += -td_not_first.pop() + d
        t_sum += 1
        ans = max(ans, d_sum + t_sum ** 2)

    return ans
    # ans = editorial_WA_TLE(N, K, td)


if __name__ == '__main__':
    main()

Submission status: Accepted

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

over_stock = []
type_set = set()

for t_i, d_i in t_and_d[:K]:
    if t_i in type_set:
        over_stock += [d_i]
    else:
        type_set.add(t_i)

x = len(type_set)
point = sum([d for t, d in t_and_d[:K]]) + x ** 2
point_max = point

for t_i, d_i in t_and_d[K:]:
    if over_stock:
        if t_i in type_set:
            continue
        else:
            diff = (x + 1) ** 2 - x ** 2 + d_i - over_stock.pop()
            type_set.add(t_i)
            x += 1
            point += diff
            point_max = max(point_max, point)

print(point_max)

Submission status: Accepted

Solution code:
n, k = [int(item) for item in input().split()]
sushi = [[int(item) for item in input().split()] for _ in range(n)]
sushi.sort(key=lambda x: x[1], reverse = True)

kind = set()
changeable = []

points = 0
for i in range(k):
  points += sushi[i][1]
  if sushi[i][0] in kind:
    changeable.append(sushi[i][1])
  else:
    kind.add(sushi[i][0])

points += len(kind)**2
ans = points

for i in range(k, n):
  if not changeable:
    break
  if sushi[i][0] not in kind:
    kind.add(sushi[i][0])
    points += sushi[i][1] - (len(kind)-1)**2 + len(kind)**2 - changeable.pop()
    ans = max(ans, points)

print(ans)

Submission status: Accepted

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


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

# pick from first
point = 0
neta_set = set()
removable_idx = list()
for i in range(K):
    s = sushi[i]
    point += s[1]
    if s[0] not in neta_set:
        neta_set.add(s[0])
    else:
        removable_idx.append(i)
point += len(neta_set) **2
max_point = point

# pick latest
for s in sushi[K:]:
    if len(removable_idx) == 0:
        break
    if s[0] in neta_set:
        continue
    idx = removable_idx[-1]
    count = len(neta_set)
    point += s[1] - sushi[idx][1]
    point += (count + 1) ** 2 - count ** 2
    neta_set.add(s[0])
    removable_idx = removable_idx[:-1]

    if point > max_point:
        max_point = point
print(max_point)

Submission status: Accepted

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
priority_que = []
t,d= [0]*N,[0]*N
sushi_max = [-1] * N
cnt = [0] * N

for i in range(N):
    t[i],d[i] = map(int,input().split(" "))
    sushi_max[t[i]-1] = max(sushi_max[t[i]-1],d[i])

t ,d= zip(*sorted(zip(t,d),reverse=True,key=lambda x: x[1]))
ans = 0
for i in range(K):
    cnt[t[i]-1] += 1
    sushi_max[t[i]-1] = -1
    pq.heappush(priority_que,(d[i],t[i]))
    ans += d[i]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz * sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(priority_que) > 0 and sushi_max[idx][0] != -1:
    minValue,minIdx = pq.heappop(priority_que)
    if cnt[minIdx-1] == 1:
        continue
    cur += ( (sz+1)*(sz+1) - (sz*sz) + sushi_max[idx][0] - minValue )
    cnt[minIdx-1] -= 1
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Accepted

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
td = []
sushi_max = [-1] * N
for i in range(N):
    t,d = map(int,input().split(" "))
    td.append((d,t))
    sushi_max[t-1] = max(sushi_max[t-1],d)

td.sort(reverse=True,key=lambda x:x[0])
cnt = [0] * N
que = []
ans = 0
for i in range(K):
    cnt[td[i][1]-1] += 1
    sushi_max[td[i][1]-1] = -1
    pq.heappush(que,td[i])
    ans += td[i][0]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz*sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(que) > 0 and sushi_max[idx][0] != -1:
    minValue,minIdx = pq.heappop(que)
    if cnt[minIdx-1] == 1:
        continue
    cur -= minValue
    cur -= (sz*sz)
    cnt[minIdx-1] -= 1
    cnt[sushi_max[idx][1]-1] += 1
    cur += (sushi_max[idx][0] + (sz+1)*(sz+1))
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(100000)


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


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


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


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


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


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


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


def mt(f):
    import time

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

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

    return wrap


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

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

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

        else:
            break
    return ans


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


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(100000)


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


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


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


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


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


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


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


def mt(f):
    import time

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

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

    return wrap


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

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

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

        else:
            break
    return max(ans)


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


if __name__ == '__main__':
    main()