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: Runtime Error

Solution code:
n, k = map(int, input().split())
ds = [[] for _ in range(n)] # ネタの種類t(<=n)別においしさdのリスト
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    ds[t].append(d)

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    # 各ネタでおいしさのsort
    y1.append(ds[i][-1])    # 各ネタで最大のおいしさをpopしy1に追加
    ds[i].pop()
    y0.extend(ds[i])        # 各ネタで2番目以降のおいしさのリストをy0に追加
y0.sort(reverse=True)   # ネタの種類が増えないおいしさのリスト（各ネタで2番目以降）を降順sort
y1.sort(reverse=True)   # ネタの種類が増えるおいしさ（各ネタで最大）のリストを降順sort

ans = 0
# Y = max(0, k-len(y0))   # Y: ネタの種類の最小値（kまで増やしながら貪欲するため）
Y = max(k, len(y1))   # Y: ネタの種類の最大値（減らしながら貪欲するため）
X = 0
for i in range(Y):
    X += y1[i]      # Y個 ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に加算
for i in range(k-Y):
    X += y0[i]      # k-Y個 ネタの種類が増えないおいしさを大きい順に加算

while True:
    ans = max(ans, X+Y*Y)
    # if (Y>=k or Y>=len(y1)):   # Yをkまで増やしながら貪欲
    #     break
    # X += y1[Y]      # ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に1つ追加
    # X -= y0[k-Y-1]  # ネタの種類が増えない大きさを小さい順に1つ除外
    # Y += 1
    if (Y<=0 or k-Y>=len(y0)):   # Yを減らしながら貪欲
        break
    X -= y1[Y-1]      # ネタの種類が増えるおいしさ（各ネタで最大）を小さい順に1つ除外
    X += y0[k-Y]  # ネタの種類が増えない大きさを大きい順に1つ追加
    Y -= 1
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
ds = [[] for _ in range(n)] # ネタの種類t(<=n)別においしさdのリスト
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    ds[t].append(d)

y0, y1 = [], []
for i in range(n):
    if len(ds[i]) == 0:
        continue
    ds[i].sort()    # 各ネタでおいしさのsort
    y1.append(ds[i][-1])    # 各ネタで最大のおいしさをpopしy1に追加
    ds[i].pop()
    y0.extend(ds[i])        # 各ネタで2番目以降のおいしさのリストをy0に追加
y0.sort(reverse=True)   # ネタの種類が増えないおいしさのリスト（各ネタで2番目以降）を降順sort
y1.sort(reverse=True)   # ネタの種類が増えるおいしさ（各ネタで最大）のリストを降順sort

ans = 0
Y = max(0, k-len(y0))   # Y: ネタの種類の最小値（kまで増やしながら貪欲するため）
X = 0
for i in range(Y):
    X += y1[i]      # Y個 ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に加算
for i in range(k-Y):
    X += y0[i]    # k-Y個 ネタの種類が増えないおいしさを大きい順に加算

while True:
    ans = max(ans, X+Y*Y)
    if (Y>=k and Y>=len(y1)):   # Yをkまで増やしながら貪欲
        break
    X += y1[Y]      # ネタの種類が増えるおいしさ（各ネタで最大）を大きい順に1つ追加
    X -= y0[k-Y-1]  # ネタの種類が増えない大きさを小さい順に1つ除外
    Y += 1
print(ans)

Submission status: Runtime Error

Solution code:
n,k=map(int,input().split())
li=[list(map(int,input().split())) for i in range(n)]
gr=[0]*(n+1) #各グループの出現回数をメモする配列

from operator import itemgetter
li=sorted(li, key=itemgetter(1), reverse=True)

A=[] #大きい方からK個以内の配列で、2個以上あるやつの2番目以降のやつの大きさを記録

for i in range(k):
  if gr[li[i][0]]>=1:
    A.append(li[i][1])
  gr[li[i][0]]+=1

A.reverse()

B=[] #既に同じグループが登場してしまってない子の大きさを記録
for i in range(k,n):
  if gr[li[i][0]]==0:
    B.append(li[i][1])
  gr[li[i][0]]+=1

#現時点での種類数とポイントを測定
sp=set()
ans=0
for i in range(k):
  ans+=li[i][1]
  sp.add(li[i][0])
sp=len(sp)
ans+=sp*sp

koukan=[] #その交換で、累積でどの程度増減するか
zouka=0
for a,b in zip(A,B):
  zouka+=(sp+1)**2-sp**2-(a-b)
  #print(a,b,sp,zouka)
  koukan.append(zouka)
  sp+=1

#koukanの中で最大の増加を探して、ansに+する。
# print(A,B)
# print(koukan)
ans+=max(0,max(koukan))

print(ans)

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
N, M = map(int, input().split())

s = [tuple(map(int, input().split())) for _ in range(N)]
s.sort(key = lambda x: -x[1])
table = [(-1, -1)] * (N + 1)
c = 0
for i in range(N):
  if table[s[i][0]] == (-1, -1):
    table[s[i][0]] = (i, s[i][1])
    c += 1

table.sort(key = lambda x: -x[1])
res = [0] * (M + 1)
vis = set()
p = M
for i in range(M): res[0] += s[i][1]
for x in range(1, c + 1):
  res[x] = res[x - 1] + pow(x, 2) - pow(x - 1, 2)
  if table[x - 1][0] >= M:
    res[x] += table[x - 1][1]
    p -= 1
    while p in vis:
      p -= 1
    res[x] -= s[p][1]
  vis.add(table[x - 1][0])
  #print(p)
print(max(res))

Submission status: Runtime Error

Solution code:
from collections import defaultdict
from itertools import accumulate
import sys
input = sys.stdin.readline

def main():
    N, K = map(int, input().split())
    d = defaultdict(list)
    for _ in range(N):
        t, v = map(int, input().split())
        d[t].append(v)

    other = []
    vmax = []
    for nums in d.values():
        vmax.append(max(nums))
        for j in nums:
            if j != max(nums):
                other.append(j)

    vmax.sort(reverse=True)
    other.sort(reverse=True)
    sum_v = list(accumulate(vmax))
    sum_o = list(accumulate(other))
    ans = 0
    for i in range(min(len(d), K)):
        tmp = (i + 1) * (i + 1) + sum_v[i]
        if K - (i + 1) > 0:
            tmp += sum_o[min(K-(i+1),len(sum_o))-1]
        ans = max(ans, tmp)

    print(ans)

    
    




if __name__ == "__main__":
    main()

Submission status: Runtime Error

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=1)

one = [[] for i in range(max(td)[0])]
point = 0
syurui = []
change = []
cnt = 0
bonus = 0
for t,d in td:
    t -= 1
    if cnt != k:
        cnt += 1
        bonus += 1
        point += d
        if t in syurui:
            bonus -= 1
            change.append([t,d])
        syurui.append(t)
    else:
        break
# print(syurui)
# print(change)
# print(point,bonus)

ans = point+bonus**2
# print(ans)
for t,d in td[k:]:
    t -= 1
    if t in syurui:
        continue
    tt,dd = change.pop()
    point -= dd
    point += d
    bonus += 1
    ans = max(ans,point+bonus**2)
    # print(change,ans)
    if not change:
        break
print(ans)

Submission status: Runtime Error

Solution code:
# D - Various Sushi

import sys

N, K = map(int, sys.stdin.buffer.readline().split())
td = []
for _ in range(N):
    tmp1, tmp2 = map(int, sys.stdin.buffer.readline().split())
    td.append([tmp1-1, tmp2, -1])

td = sorted(td, key = lambda x: x[1], reverse = True)
t_counter = [0]*N
for idx in range(N):
    t_counter[td[idx][0]] += 1
    td[idx][2] = t_counter[td[idx][0]]

t_set = set([td[idx][0] for idx in range(K)])
ans_list = []
ans = sum([td[idx][1] for idx in range(K)]) + len(t_set)**2
ans_list.append(ans)

a = 0
b = 0
for add_kind in range(1, len([x for x in t_counter if x > 0]) - len(t_set) + 1):
    while 1:
        if td[K - 1 - a][2] >= 2:
            ans -= td[K - 1 - a][1]
            a += 1
            break
        a += 1
    while 1:
        if td[K + b][0] not in t_set:
            ans -= len(t_set)**2
            t_set.add(td[K + b][0])
            ans += td[K + b][1] + len(t_set)**2
            b += 1
            break
        b += 1
    ans_list.append(ans)

print(max(ans_list))

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter
from heapq import heapify, heappop, heappush
read = sys.stdin.read

N, K, *td = map(int, read().split())
dt = sorted(zip(td[1::2], td[::2]), reverse=True)

candidates = []
neta_group = Counter([j for i, j in dt[:K]])
variety = len(neta_group.keys())
points = sum([i for i, j in dt[:K]]) + variety ** 2
candidates.append(points)

if N == K:
    print(points)
    exit()

heap = [(-i, j) for i, j in dt[K:]]
heapify(heap)
satisfaction = dt[:K]
heapify(satisfaction)

while heap:
    taste, neta = heappop(heap)
    taste = -taste
    if not neta_group.get(neta):
        taste2, neta2 = heappop(satisfaction)
        if neta_group[neta2] == 1:
            heappush(heap, (-taste2, neta2))
        else:
            neta_group[neta2] -= 1
            neta_group[neta] = 1
            heappush(heap, (-taste, neta))
            points = points - taste2 + taste - variety ** 2 + (variety + 1) ** 2

            variety += 1
            candidates.append(points)

print(max(candidates))

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter
from heapq import heapify, heappop, heappush
read = sys.stdin.read

N, K, *td = map(int, read().split())
dt = sorted(zip(td[1::2], td[::2]), reverse=True)

candidates = []
neta_group = Counter([j for i, j in dt[:K]])
number = len(neta_group)
points = sum([i for i, j in dt[:K]]) + number ** 2
candidates.append(points)

heap = [(-i, j) for i, j in dt[K:]]
heapify(heap)
manzoku = dt[:K]
heapify(manzoku)

if N != K:
    while heap:
        oishisa, neta = heappop(heap)
        oishisa = -oishisa
        if not neta_group.get(neta):
            oishisa2, neta2 = heappop(manzoku)
            if neta_group[neta2] == 1:
                heappush(heap, (-oishisa2, neta2))
            else:
                neta_group[neta2] -= 1
                neta_group[neta] = 1
                heappush(heap, (-oishisa, neta))
                points = points - oishisa2 + oishisa - number ** 2 + (number + 1) ** 2

                number += 1
                candidates.append(points)

print(max(candidates))

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter
from heapq import heapify, heappop, heappush
read = sys.stdin.read

N, K, *td = map(int, read().split())
dt = sorted(zip(td[1::2], td[::2]), reverse=True)

candidates = []
neta_group = Counter([j for i, j in dt[:K]])
number = len(neta_group)
points = sum([i for i, j in dt[:K]]) + number ** 2
candidates.append(points)

heap = [(-i, j) for i, j in dt[K:]]
heapify(heap)
manzoku = dt[:K]
heapify(manzoku)

while heap:
    oishisa, neta = heappop(heap)
    oishisa = -oishisa
    if not neta_group.get(neta):
        oishisa2, neta2 = heappop(manzoku)
        if neta_group[neta2] == 1:
            heappush(heap, (-oishisa2, neta2))
        else:
            neta_group[neta2] -= 1
            neta_group[neta] = 1
            heappush(heap, (-oishisa, neta))
            points = points - oishisa2 + oishisa - number ** 2 + (number + 1) ** 2

            number += 1
            candidates.append(points)

print(max(candidates))

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter
from heapq import heapify, heappop, heappush
read = sys.stdin.read

N, K, *td = map(int, read().split())
dt = sorted(zip(td[1::2], td[::2]), reverse=True)

candidates = []
points = sum([i for i, j in dt[:K]]) + len(set([j for i, j in dt[:K]])) ** 2
neta_group = Counter([j for i, j in dt[:K]])
candidates.append(points)

heap = [(-i, j) for i, j in dt[K:]]
heapify(heap)
manzoku = dt[:K]
heapify(manzoku)

while heap:
    oishisa, neta = heappop(heap)
    oishisa = -oishisa
    if not neta_group.get(neta):
        oishisa2, neta2 = heappop(manzoku)
        if neta_group[neta2] == 1:
            heappush(heap, (-oishisa2, neta2))
        else:
            neta_group[neta2] -= 1
            neta_group[neta] = 1
            heappush(heap, (-oishisa, neta))
            points = points - oishisa2 + oishisa - (len(neta_group.keys()) - 1) ** 2 + len(neta_group.keys()) ** 2
            candidates.append(points)

print(max(candidates))

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(100000)
n, k = list(map(int, input().split(' ')))
sushi = {}
for i in range(n):
    t, d = list(map(int, input().split(' ')))
    sushi[t] = sushi.get(t, []) + [d]
for key in sushi:
    sushi[key] = sorted(sushi[key])[::-1]
# print(sushi)

memo = {}
used = [False] * (len(sushi) + 1)
def eat(sushi, ct, pt):
    hoge = str(sorted(sushi.items()))
    if hoge in memo:
        return memo[hoge]

    # 終わり
    if ct == k:
        # print (pt, sum([v for _, v in used.items()]) ** 2)
        return pt + sum(used) ** 2
    same = 0
    diff = 0

    # 同じやつ
    target = None
    max_val = -1
    for key in sushi:
        if used[key]:
            if sushi[key] and max_val < sushi[key][0]:
                target = key
                max_val = sushi[key][0]
    if target is not None:
        val = sushi[target].pop(0)
        same = eat(sushi, ct+1, pt+val)
        # 戻す
        sushi[target].insert(0, val)

    # 新しいやつ
    target = None
    max_val = -1
    for key in sushi:
        if not used[key]:
            if sushi[key] and max_val < sushi[key][0]:
                target = key
                max_val = sushi[key][0]
    if target is not None:
        val = sushi[target].pop(0)
        used[target] = True
        diff = eat(sushi, ct+1, pt+val)
        # 戻す
        used[target] = False
        sushi[target].insert(0, val)
    memo[hoge] = max(same, diff)
    return memo[hoge]

print (eat(sushi, 0, 0))

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
td = list(sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1]))
td2 = []
eaten = set()
dup = 0
score = k * k
i = 0
j = 0
while len(eaten) + dup < k:
    if td[i][0] in eaten:
        td2.append(td[i])
        i += 1
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if i < n and (j >= len(td2) or td[i][1] >= td2[j][1] - diff):
            t, d = td[i]
            score += d
            eaten.add(t)
            i += 1
        else:
            t, d = td2[j]
            score += d - diff
            dup += 1
            j += 1
print(score)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
td = list(sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1]))
td2 = []
eaten = set()
dup = 0
score = k * k
while len(eaten) + dup < k:
    if td[0][0] in eaten:
        td2.append(td.pop(0))
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if not td2 or td[0][1] >= td2[0][1] - diff:
            t, d = td.pop(0)
            score += d
            eaten.add(t)
        else:
            t, d = td2.pop(0)
            score += d - diff
            dup += 1

print(score)

Submission status: Runtime Error

Solution code:
from heapq import heappush, heappop
N, K = map(int, input().split())
log = [0]*N
pre = []
ls = []
for i in range(N):
  t, d = map(int, input().split())
  t = t-1
  heappush(ls,(-d,t))
v = 0
n = 0
for i in range(K):
  d, t = heappop(ls)
  d = -d
  log[t] += 1
  v += d
  if log[t]>1:
    heappush(pre,(d,t))
  else:
    n += 1
v += n*n
ans = v
while ls:
  d, t = heappop(ls)
  if log[t]>0:
    continue
  log[t] += 1
  d = -d
  pv = v
  x, y = heappop(pre)
  v -= x
  v += d
  v += 2*n+1
  n += 1
  if ans<v:
    ans = v
print(ans)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(2147483647)
INF=float("inf")
MOD=10**9+7
input=lambda :sys.stdin.readline().rstrip()
def resolve():
    from heapq import heappop,heappush,heapify
    n,k=map(int,input().split())
    A=[0]*n
    for i in range(n):
        t,d=map(int,input().split())
        t-=1
        A[i]=(t,d)
    A.sort(lambda x:-x[1]) # d の降順ソート

    # 各種の中で最大かどうかをわける
    B=[0]*n
    first=[1]*n
    for i,a in enumerate(A):
        t,d=a
        if(first[t]):
            first[t]=0
            B[i]=(d,1)
        else:
            B[i]=(d,0)

    # 初めは貪欲に d の大きい方から k 個取る
    S=[]
    for i in range(k):
        d,f=B[i]
        heappush(S,(f,d))
    x=sum(S[i][1] for i in range(k))
    y=sum(S[i][0] for i in range(k))
    ans=x+y*y

    # 残りの n-k について、Sからf=0で最小のものを捨て、B[k:]からf=1の中で最大のものを加える
    B=B[k:]
    for i in range(n-k):
        d,f=B[i]
        if(f==0): continue
        pf,pd=heappop(S)
        if(pf==1): break
        x+=d-pd
        y+=1
        ans=max(ans,x+y*y)
    print(ans)
resolve()

Submission status: Runtime Error

Solution code:
import heapq
N, K = [int(x) for x in input().split()]
TD = [[int(x) for x in input().split()] for _ in range(N)]
kinds = set()
h1 = []
h2 = []
for t, d in sorted(TD, key=lambda td: td[1], reverse=True):
    if (t in kinds):
        heapq.heappush(h1, [-d, 0])
    else:
        kinds.add(t)
        heapq.heappush(h1, [-d, 1])
X = 0
Y = 0
ans = 0
for i in range(K):
    xi, yi = heapq.heappop(h1)
    X += -xi
    Y += yi
    heapq.heappush(h2, [yi, -xi])
ans = [X + Y*Y]
while(len(h1)):
    if (Y == len(kinds)): break
    xi, yi = heapq.heappop(h1)
    if (yi == 0): continue
    Y += 1
    yj, xj = heapq.heappop(h2)
    X += (-xj) + (-xi)
    ans.append(X + Y*Y)
print(max(ans))

Submission status: Runtime Error

Solution code:
# 同じ種類だと、点数が高いやつを選ぶ
# 違う種類を選ぶと、ボーナス点が入るので、それをどうするか
# d  = sum(di) + x**x
# 種類xを固定すると？ O(log(N))でいける？
# ※※ 種類が増えるかどうかを降順で見ながらboolで置き換えるという発想。強い。

from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []
# まずは上からK個貪欲に取る
for i in range(k):
    d,t = sushi[i]
    heapq.heappush(sushiq,(d,t)) 
    sushidict[t] += 1
    cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(k,n):
    d,t = sushi[i]
    # すでにある種類の場合は大丈夫
    if sushidict[t] > 0 : continue
    # 種類が増える時、小さいものから入れ替えていく
    d_min,t_min = heapq.heappop(sushiq)
    # popして消えるようなものはNG
    while sushidict[t_min] == 1:
        d_min, t_min = heapq.heappop(sushiq)
    sushidict[t_min] -= 1
    sushidict[t] += 1
    n_kind += 1
    cnt  = cnt - d_min + d
    res.append(cnt+n_kind**2)
    if n_kind == n_kind_all:
        break

print(max(res))
# print(res)

Submission status: Runtime Error

Solution code:
# 同じ種類だと、点数が高いやつを選ぶ
# 違う種類を選ぶと、ボーナス点が入るので、それをどうするか
# d  = sum(di) + x**x
# 種類xを固定すると？ O(log(N))でいける？
# ※※ 種類が増えるかどうかを降順で見ながらboolで置き換えるという発想。強い。

from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []
# まずは上からK個貪欲に取る
for i in range(k):
    heapq.heappush(sushiq,(d,t)) 
    d,t = sushi[i]
    sushidict[t] += 1
    cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(k,n):
    d,t = sushi[i]
    # すでにある種類の場合は大丈夫
    if sushidict[t] > 0 : continue
    # 種類が増える時、小さいものから入れ替えていく
    d_min,t_min = heapq.heappop(sushiq)
    # popして消えるようなものはダメ
    while sushidict[t_min] == 0:
        d_min, t_min = heapq.heappop(sushiq)
        print(d_min,t_min)
    sushidict[t_min] -= 1
    sushidict[t] += 1
    n_kind += 1
    cnt  = cnt - d_min + d
    res.append(cnt+n_kind**2)
    if n_kind == n_kind_all:
        break

print(max(res))

Submission status: Runtime Error

Solution code:
import sys
import heapq
from copy import deepcopy
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 7)


def main():
    n, k = map(int, input().split())
    td = []
    for _ in range(n):
        t, d = map(int, input().split())
        td.append([-d, t])
    td2 = deepcopy(td)
    heapq.heapify(td)
    heapq.heapify(td2)

    ans = 0
    s = set()
    for _ in range(k):
        tmp = heapq.heappop(td)
        ans += -tmp[0]
        s.add(tmp[1])
    l = len(s)

    ans2 = 0
    s2 = set()
    while len(s2) != k:
        tmp = heapq.heappop(td2)
        if tmp[1] in s2:
            continue
        ans2 += -tmp[0]
        s2.add(tmp[1])

    l2 = len(s2)

    print(max((ans + l * l), (ans2 + l2 * l2)))


if __name__ == '__main__':
    main()

Submission status: Runtime Error

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

l=[] # 与えられた寿司のリスト。寿司をタプル(ネタti,おいしさdi)で表す。既に食べた種類の寿司は，そのおいしさdiをdjに移す。食べた寿司はリストから削除する。
dj=[] # 既に食べた種類の寿司の「おいしさ基礎ポイント」を格納するリスト。食べた寿司の「おいしさ基礎ポイント」はリストから削除する。
s=set() # 既に食べた寿司の種類を格納するセット

for i in range(N):
    t,d=map(int,input().split())
    l.append((t,d))
l.sort(key=lambda x:x[1],reverse=True) # l内の寿司をおいしさdiで降順にソート

ans=K**2
j=0 # 既に食べたことのある種類の寿司を更に食べた回数

for i in range(K): # i番目に食べる寿司について
    while l and l[0][0] in s: # lが空でなく，先頭の寿司の種類がsに入っている（先頭の寿司が食べたことのある種類である）場合
        dj.append((l.pop(0))[1]) # 食べたことのない種類の寿司がlの先頭に現れるか，lが空になるまで，先頭の寿司のおいしさをdjに格納する
    if not dj or l[0][1]>=dj[0]-(2*K-1-2*j): # djが空である場合，又は，lの先頭の寿司のおいしさと，djの先頭の「寿司のおいしさ」から「種類ボーナスポイント」の差分を減じたものを比較し，前者が勝る場合
        ans+=l[0][1] # lの先頭の寿司のおいしさに「種類ボーナスポイント」の差分を加えたものを答えに加算
        s.add(l[0][0]) # その寿司の種類をsに追加
        l.pop(0) # その寿司をlから削除
    else: #　後者が勝る場合
        ans+=dj[0]-(2*K-1-2*j) # djの先頭の「寿司のおいしさ」を答えに加算
        dj.pop(0) # その「寿司のおいしさ」をdjから削除
        j+=1

print(ans)

Submission status: Runtime Error

Solution code:
from heapq import heappush, heappop

INF = float("inf")

N, K, *TD = map(int, open(0).read().split())

S = [[] for _ in range(N)]
for t, d in zip(*[iter(TD)] * 2):
    S[t - 1].append(d)

E = []
for s in S:
    if s:
        E.append(sorted(s, reverse=True))
    else:
        E.append([-INF])
E.sort(reverse=True)

cur = 0
Q = []
for i in range(K):
    cur += E[i][0]
    for j in range(1, len(E[i])):
        heappush(Q, -E[i][j])

for i in range(K, N):
    for j in range(len(E[i])):
        heappush(Q, -E[i][j])

res = cur + K * K
for x in range(K, 0, -1):
    v, w = E[x][0], -Q[0]
    if v < w:
        heappop(Q)
        cur += w
        heappush(Q, -v)
        cur -= v

    res = max(res, cur + x * x)

print(res)

Submission status: Runtime Error

Solution code:
import heapq
from collections import deque
n,k=map(int,input().split())
l=[list(map(int,input().split())) for i in range(n)]

for i in range(n):
  l[i].reverse()
  l[i][0]*=-1
heapq.heapify(l)

check=[0 for i in range(10**5+1)]
L=[]
L2=[]
L2=deque(L2)
ans=[]

for i in range(k):
  x=heapq.heappop(l)
  if check[x[1]]==0:
    L.append(-x[0])
    check[x[1]]=1
  else:
    L2.append(-x[0])
p=sum(check)
ans.append(sum(L)+sum(L2)+p**2)

y=n-k
while y!=0:
  x=heapq.heappop(l)
  if check[x[1]]==0:
    L.append(-x[0])
    check[x[1]]=1
    z=L2.pop()
    p+=1
    ans.append(ans[-1]-(p-1)**2-x[0]-z+p**2)
  y-=1

print(max(ans))

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
from collections import defaultdict
from collections import deque
from heapq import heappush, heappop
import sys
import math
import bisect
import random
import itertools
sys.setrecursionlimit(10**5)
stdin = sys.stdin
bisect_left = bisect.bisect_left
bisect_right = bisect.bisect_right
def LI(): return list(map(int, stdin.readline().split()))
def LF(): return list(map(float, stdin.readline().split()))
def LI_(): return list(map(lambda x: int(x)-1, stdin.readline().split()))
def II(): return int(stdin.readline())
def IF(): return float(stdin.readline())
def LS(): return list(map(list, stdin.readline().split()))
def S(): return list(stdin.readline().rstrip())
def IR(n): return [II() for _ in range(n)]
def LIR(n): return [LI() for _ in range(n)]
def FR(n): return [IF() for _ in range(n)]
def LFR(n): return [LI() for _ in range(n)]
def LIR_(n): return [LI_() for _ in range(n)]
def SR(n): return [S() for _ in range(n)]
def LSR(n): return [LS() for _ in range(n)]
mod = 1000000007
inf = float('INF')

#A
def A():
    a = LI()
    a.sort()
    print(a[0]*a[1]//2)  
    return

#B
def B():
    d = defaultdict(int)
    s = II()
    def f(n):
        if n % 2:
            return 3 * n + 1
        return n // 2
    i = 1
    while not d[s]:
        d[s] = 1
        s = f(s)
        i += 1
    print(i)
    return

#C
def C():
    n = II()
    h = LI()
    ans = h[0]
    for i in range(n - 1):
        ans += max(0, h[i + 1] - h[i])
    print(ans)
    return

# D
# 解説AC
# 大きいものからとっていった際にその種類数をn
# としてi<=nの種類数においてこれを超えるものは無い
# それより大きいものについて種類数が減らずかつとったものの中で
# 一番小さいものを除き、種類数が増える一番大きいものをとってこれば
# i+1番目の種類下でもっとも大きい数になる 
def D():
    n, k = LI()
    ma = []
    for _ in range(n):
        t, di = LI()
        ma.append((di, t))
    ma.sort()
    ma = deque(ma)
    s = defaultdict(int)
    t = 0
    ans = 0
    q = []
    for _ in range(k):
        di,ti = ma.pop()
        ans += di
        heappush(q, (di, ti))
        if not s[ti]:
            t += 1
        s[ti] += 1
    ans += t ** 2
    ans = [ans]
    while ma and q:
        di, ti = ma.pop()
        while s[ti] and ma:
            di, ti = ma.pop()
        if not ma and s[ti]:
            break
        d1, t1 = heappop(q)
        while s[t1] == 1 and q:
            d1, t1 = heappop(q)
        if not q and s[t1]] == 1:
            break
        s[t1] -= 1
        tmp = ans[-1]
        tmp = tmp - d1 + di - t ** 2
        t += 1
        tmp += t ** 2
        ans.append(tmp)
        heappush(q, (di, ti))
        s[ti] += 1
    print(max(ans))


    return

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

Submission status: Runtime Error

Solution code:
from operator import itemgetter

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

s = []

last_s = 0
for i in range(n):
    t, d = map(int, input().split())
    s.append((t, d))

ss = sorted(s, key=itemgetter(1), reverse=True)



def calc(l, num, cnt):
    if num == k or cnt == n:
        pts = 0
        kinds = set()

        for li in l:
            if li[0] not in kinds:
                kinds.add(li[0])
            pts += li[1]
        pts += len(kinds) ** 2
        # print(str(l) + ' ' + str(num) + ' ' + str(cnt)+' '+str(pts)+' '+str(kinds)+' '+str(len(kinds)))
        return pts
    return max(calc(l + [ss[cnt]], num + 1, cnt + 1), calc(l, num, cnt + 1))


print(calc([], 0, 0))

Submission status: Runtime Error

Solution code:
from heapq import heappop, heappush
N, K = map(int, input().split())
KMAX = 10 ** 5 + 1
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((d, t))

TD.sort(reverse=True)

pioneer = []
successor = []
appear = [False] * KMAX
cnt = 0
nsp = 0  # num of species
base = 0
best = []

for td in TD:
    d, t = td
    if cnt < K:
        heappush(successor if appear[t] else pioneer, (d, t))
        nsp += 0 if appear[t] else 1
        appear[t] = True
        base += d
        cnt += 1
        if cnt == K-1:
            heappush(best, -(base + nsp ** 2))
        continue
    if not successor:
        break
    if appear[t]:
        continue
    cd, ct = heappop(successor)
    appear[t] = True
    heappush(pioneer, (d, t))
    base += (d-cd)
    nsp += 1
    heappush(best, -(base + nsp ** 2))

print(-heappop(best))

Submission status: Runtime Error

Solution code:
from heapq import heappop, heappush
N, K = map(int, input().split())
KMAX = 10 ** 5 + 1
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((d, t))

TD.sort(reverse=True)

pioneer = []
successor = []
appear = [False] * KMAX
cnt = 0
nsp = 0  # num of species
base = 0
best = []

for td in TD:
    d, t = td
    if cnt < K:
        heappush(successor if appear[t] else pioneer, (d, t))
        nsp += 0 if appear[t] else 1
        appear[t] = True
        base += d
        cnt += 1
        continue
    elif cnt == K:
        heappush(best, -(base + nsp ** 2))
    if not successor:
        break
    if appear[t]:
        continue
    cd, ct = heappop(successor)
    appear[t] = True
    heappush(pioneer, (d, t))
    base += (d-cd)
    nsp += 1
    heappush(best, -(base + nsp ** 2))

print(-heappop(best))

Submission status: Runtime Error

Solution code:
n, m = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for n in range(n)]
td.sort(key=lambda x:-x[1])
memo = set()
a = []
for t, d in td:
  if t in memo:
    a.append((d, 0))
  else:
    a.append((d, 1))
  memo.add(t)
a = [(-x, x, d) for x, d in a]
import heapq
heapq.heapify(a)
val = 0
kind = 0
b = []
for _ in range(m):
  ele = heapq.heappop(a)
  val += ele[1]
  kind += ele[2]
  if ele[2] == 0:
    b.append(ele[1])
ans = val + kind ** 2
while (len(a) > 0):
  val -= b.pop()
  flag = False
  while(len(a) > 0):
    elem = heapq.heappop(a)
    if elem[2] == 1:
      flag = True
      break
  if not flag:
    break
  val += elem[1]
  kind += 1
  tmpans = val + kind ** 2
  if tmpans > ans:
    ans = tmpans
print(ans)

Submission status: Runtime Error

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

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2
huku.sort(key = lambda x:x[1], reverse = True)

ansl = []
ansl.append(ans)
for i,j in TD[k:]:
  if i in su:
    continue
  elif:
    huku == []
    break
  else:
    temp = ansl[-1] + j - huku.pop()[1] + (len(su)+1)**2 -len(su)**2
    su.add(i)
    ansl.append(temp)

print(max(ansl))

Submission status: Runtime Error

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

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2
huku.sort(key = lambda x:x[1], reverse = True)

ansl = []
ansl.append(ans)
for i,j in TD[k:]:
  if i in su:
    continue
  else:
    temp = ansl[-1] + j - huku.pop()[1] + (len(su)+1)**2 -len(su)**2
    su.add(i)
    ansl.append(temp)

print(max(ansl))

Submission status: Runtime Error

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

ans = 0
huku = []
su = set([])

for i, j in TD[:k]:
  ans += j
  if i in su:
    huku.append([i,j])
  su.add(i)
ans += len(su)**2

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

for i,j in TD[k:]:
  if i in su:
    continue
  else:
    su.add(i)
    temp = ans + j - huku.pop()[1] + len(su)**2 -(len(su)-1)**2
    ans = max(ans, temp)

print(ans)

Submission status: Runtime Error

Solution code:
mi = lambda: map(int, input().split())
n, k = mi()
sushis = [list(mi()) for _ in range(n)]
types = set()
sushis.sort(key = lambda x: x[1], reverse = True)
bp = tp = sp = 0
exts = []
for t, d in sushis[:k]:
    bp += d
    if t in types:
        exts += d,
    else:
        types.add(t)
        tp += 1
sp = bp+tp**2
for t, d in sushis[k:]:
    if t not in types:
        types.add(t)
        bp += d-exts.pop()
        tp += 1
        sp = max(sp, bp+tp**2)
print(sp)