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:
import heapq as hq

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

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

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

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

Submission status: Runtime Error

Solution code:
n,k= map(int,input().split())
td=[]
for i in range(n):
    td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
select=td[:k]
cat=list(set([tdi[0] for tdi in select]))
noselect=list(filter(lambda x: x[0] not in cat,td[k:]))
noselect.sort(key=lambda x:x[0])
noselect02=[]
noselect02.append(noselect[0])
for no in noselect[1:]:
    if no[0] != noselect[-1][0]:
        noselect.append(no)
        continue
    if no[1] > noselect[-1][1]:
        noselect[1]=no[1]

rs = sum([tdi[1] for tdi in select])+len(set([tdi[0] for tdi in select]))**2
for i in range(1,min(k,len(noselect)+1)):
    t = select[:-i]
    t[k-i:k-i] = noselect[:i]
    r = sum([tdi[1] for tdi in t])
    r += len(set([tdi[0] for tdi in t]))**2
    if rs < r:
        rs = r
print(rs)

Submission status: Runtime Error

Solution code:
n,k= map(int,input().split())
td=[]
for i in range(n):
    td.append(list(map(int,input().split())))
td.sort(key=lambda x:x[1],reverse=True)
select=td[:k]
cat=list(set([tdi[0] for tdi in select]))
noselect=[]
tdk=td[k:]
tdk.sort(key=lambda x:x[0])
noselect.append(tdk[0])
for no in tdk[1:]:
    if no[0] != noselect[-1][0]:
        noselect.append(no)
        continue
    if no[1] > noselect[-1][1]:
        noselect[1]=no[1]

rs = sum([tdi[1] for tdi in select])+len(set([tdi[0] for tdi in select]))**2
for i in range(1,min(k,len(noselect)+1)):
    t = select[:-i]
    t[k-i:k-i] = noselect[:i]
    r = sum([tdi[1] for tdi in t])
    r += len(set([tdi[0] for tdi in t]))**2
    if rs < r:
        rs = r
print(rs)

Submission status: Runtime Error

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
top = [0] * 10 ** 5
sub = []

for i in range(n):
    t, d = map(int, input().split())
    if top[t] == 0:
        top[t] = d
    else:
        if d > top[t]:
            sub.append(top[t])
            top[t] = d
        else:
            sub.append(d)

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

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

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

ans = 0

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

Submission status: Runtime Error

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
top = [0] * 10 ** 5
sub = []
kinds = 0

for i in range(n):
    t, d = map(int, input().split())
    if top[t] == 0:
        top[t] = d
        kinds += 1
    else:
        if d > top[t]:
            sub.append(top[t])
            top[t] = d
        else:
            sub.append(d)

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

for p in top:
    top_sum.append(top_sum[-1] + p)

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

ans = 0
start = min(k, len(top_sum) - 1)
end = max(1, k - len(sub_sum) + 1)

for i in range(start, end - 1, -1):
    ans = max(ans, top_sum[i] + sub_sum[k - i] + i ** 2)
print(ans)

Submission status: Runtime Error

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 t in s:
            continue
        m += d - rep.pop()
        s.add(t)
        p = max(p, len(s) ** 2 + m)
        if not rep:
            break
    return p

print(main())

Submission status: Runtime Error

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

cnt = 0
while True:
    allzero = True
    for i in range(N):
        if(h[i] != 0):
            allzero = False
    if allzero: break

    l = 0
    while h[l] == 0:l += 1
    r = l
    while r <N-1 and h[r+1]!=0:
        r +=1
    cnt += 1
    for i in range(l,r + 1,1):
        h[i] -=1
print(cnt)

Submission status: Runtime Error

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

Solution code:
def Main():
    N, K = map(int, input().split())
    Sushi = []
    for i in range(N):
        t, d = map(int, input().split())
        Sushi.append((t, d))

    memo = {}
    for t, d in Sushi:
        if t not in memo:
            memo[t] = d
        else:
            if memo[t] < d:
                memo[t] = d
    zero = []
    one = []
    for t, d in Sushi:
        if t in memo and memo[t] == d:
            one.append(d)
            del memo[t]
        else:
            zero.append(d)

    zero.sort(reverse=True)
    one.sort(reverse=True)

    neta = len(one)

    two = one + zero
    cx = sum(two[:K])
    if neta >= K:
        max_hp = cx + K**2
    else:
        max_hp = cx + neta**2

    current_zero = 0 if neta <= K else K - neta
    current_one = K-1 if neta >= K else neta-1
    while current_one >= 0 and current_zero < len(zero):
        cx = cx - one[current_one] + zero[current_zero]
        hp = cx + current_one**2
        if hp > max_hp:
            max_hp = hp
        current_zero += 1
        current_one -= 1

    print(max_hp)


if __name__ == '__main__':
    Main()

Submission status: Runtime Error

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

sat = sum(used)**2 + deli

dup.sort(reverse=True)

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

print(sat)

Submission status: Runtime Error

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):
        candel.append(heapq.heappop(sumn[k-1]))
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: Runtime Error

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]
    if upperKCount < K:
        for k in range(upperKCount + 1, K+1):
            maxPoint = max(maxPoint, k ** 2 + typeMaxSum[k] + typeLowerSum[K - k])
    print(maxPoint)
    return 0

if __name__ == "__main__":
    solve()

Submission status: Runtime Error

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, K+1):
        maxPoint = max(maxPoint, k ** 2 + typeMaxSum[k] + typeLowerSum[K - k])
    print(maxPoint)
    return 0

if __name__ == "__main__":
    solve()

Submission status: Runtime Error

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(K-len(Y),min(len(SX),K+1)):
    ans=max(SX[x]+SY[K-x]+x*x,ans)
print(ans)

Submission status: Runtime Error

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 or len(unuse) != 0 or 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: Runtime Error

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

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

q = []
for _ in range(N):
    t, d = map(int, input().split())
    heappush(q, (-d, t))

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
for _ in range(K-used):
    d, t = heappop(not_used)
    su -= d
    heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    heappush(now_used, (-a, grp))
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

Submission status: Runtime Error

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

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

q = []
for _ in range(N):
    t, d = map(int, input().split())
    heappush(q, (-d, t))

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
for _ in range(K-used):
    d, t = heappop(not_used)
    su -= d
    heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

Submission status: Runtime Error

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

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

q = []
for _ in range(N):
    t, d = map(int, input().split())
    heappush(q, (-d, t))

now_used = []
not_used = []

c = 0
su = 0
kn = set()
while q and c<K:
    d, t = heappop(q)
    if not t in kn:
        kn.add(t)
        heappush(now_used, (-d, t))
        su -= d
        c += 1
    else:
        heappush(not_used, (d, t))

used = len(now_used)
if K-used:
    for _ in range(K-used):
        d, t = heappop(not_used)
        su -= d
        heappush(now_used, (-d, t))

ans = su+c*c

while not_used:
    a, grp = heappop(now_used)
    su -= a
    kn.remove(grp)
    a, grp = heappop(not_used)
    kn.add(grp)
    su -= a
    c = len(kn)
    ans = max(ans, su+c*c)
print(ans)

Submission status: Runtime Error

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

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)
  if len(candidates) == 0:
    break

print(max(ans))

Submission status: Runtime Error

Solution code:
n,k=map(int,input().split())
l=[]
for i in range(n):
    l.append(list(map(int,input().split())))
p=[[] for i in range(n)]
for i in range(n):
    p[l[i][0]-1].append(l[i][1])
a=[]
b=[]
for i in range(n):
    if len(p[i])>0:
        mp=max(p[i])
        f=0
        for j in p:
            if f==0 and j==mp:
                a.append(j)
                f+=1
            else:
                b.append(j)
ans=len(a)*len(a)+sum(a)
if b:
    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: Runtime Error

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

list0 = []
list1 = []
set1 = set()
for i in range(N):
    t, d = map(int, input().split())
    if t in set1:
        list0.append(d)
    else:
        set1.add(t)
        list1.append(d)

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

various = max(K, len(list1))
selected = []
for i in range(various):
    selected.append(list1.pop(0))

ans = sum(selected) + various ** 2

while list0 and list0[0] > selected[-1]:
    selected.pop(-1)
    selected.insert(0, list0.pop(0))
    various -= 1
    tmp = sum(selected) + various ** 2
    if tmp <= ans:
        break
    ans = tmp

print(ans)

Submission status: Runtime Error

Solution code:
import heapq
import math

N,K =  list(map(int,input().split(" ")))
neta = set()
neta_dict = dict()
sushis = [list(map(int,input().split(" "))) for _ in range(N)]
sushis = sorted(sushis,key = lambda x: -x[1])
# print(sushis)
ans = 0
種類 = 0

for i in range(K): # とりあえず美味しさで最適化してみる
    k,v = sushis[i]
    neta.add(k)
    if k not in neta_dict:
        neta_dict[k] = [v]
        種類 += 1
    else:
        neta_dict[k].append(v)
    ans += v
ans += (種類) ** 2
final_ans = ans
que= []
for i in range(K,N): # 種類を最適化する際に使う寿司たち
    heapq.heappush(que,(-sushis[i][1],sushis[i][0]  )) # 美味しさでpopするため順番を入れ替える

nums = [] # ここに美味しさを犠牲にする際に捨てる寿司を打ち込む
for k,v in neta_dict.items():
    if len(v) > 1:
        v.sort()
        for temp in v[:-1]:
            nums.append(temp)
nums.sort() # 少ない順に捨てていく
nums_i = 0
while que:
    v,k = heapq.heappop(que)
    v = -v
    if k not in  neta:
        neta.add(k) # 追加
        ans -= nums[nums_i]
        ans += v
        ans -= 種類 * 種類
        種類 += 1
        ans += 種類 * 種類
        final_ans = max(final_ans,ans)

print(final_ans)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
import os
import itertools
import string
import heapq
import _collections
from collections import Counter
from collections import defaultdict
from functools import lru_cache
import bisect
import re
import queue


class Scanner():
    @staticmethod
    def int():
        return int(sys.stdin.readline().rstrip())

    @staticmethod
    def string():
        return sys.stdin.readline().rstrip()

    @staticmethod
    def map_int():
        return [int(x) for x in Scanner.string().split()]

    @staticmethod
    def string_list(n):
        return [input() for i in range(n)]

    @staticmethod
    def int_list_list(n):
        return [Scanner.map_int() for i in range(n)]

    @staticmethod
    def int_cols_list(n):
        return [int(input()) for i in range(n)]


class Math():
    @staticmethod
    def gcd(a, b):
        if b == 0:
            return a
        return Math.gcd(b, a % b)

    @staticmethod
    def lcm(a, b):
        return (a * b) // Math.gcd(a, b)

    @staticmethod
    def roundUp(a, b):
        return -(-a // b)

    @staticmethod
    def toUpperMultiple(a, x):
        return Math.roundUp(a, x) * x

    @staticmethod
    def toLowerMultiple(a, x):
        return (a // x) * x

    @staticmethod
    def nearPow2(n):
        if n <= 0:
            return 0
        if n & (n - 1) == 0:
            return n
        ret = 1
        while(n > 0):
            ret <<= 1
            n >>= 1
        return ret

    @staticmethod
    def sign(n):
        if n == 0:
            return 0
        if n < 0:
            return -1
        return 1

    @staticmethod
    def isPrime(n):
        if n < 2:
            return False
        if n == 2:
            return True
        if n % 2 == 0:
            return False
        d = int(n ** 0.5) + 1
        for i in range(3, d + 1, 2):
            if n % i == 0:
                return False
        return True


MOD = int(1e09) + 7
INF = int(1e15)

T_MAX = int(1e05) + 1


def main():
    # sys.stdin = open("sample.txt")
    N, K = Scanner.map_int()
    T = defaultdict(list)
    for _ in range(N):
        t, d = Scanner.map_int()
        t -= 1
        T[t].append(d)
    for v in T.values():
        v.sort(reverse=True)
    M = []
    for v in T.values():
        M.append((v[0], 1))
        for i in range(1, len(v)):
            M.append((v[i], 0))
    M.sort(reverse=True)
    q = []
    x = 0
    y = 0
    for i in range(K):
        x += M[i][0]
        if M[i][1] == 0:
            heapq.heappush(q, M[i][0])
        y += M[i][1]
    ans = x + y * y
    for i in range(K, N):
        if M[i][1] == 1:
            x -= q.pop()
            x += M[i][0]
            y += 1
            ans = max(ans, x + y * y)
            if q == []:
                break
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
import os
import itertools
import string
import heapq
import _collections
from collections import Counter
from collections import defaultdict
from functools import lru_cache
import bisect
import re
import queue


class Scanner():
    @staticmethod
    def int():
        return int(sys.stdin.readline().rstrip())

    @staticmethod
    def string():
        return sys.stdin.readline().rstrip()

    @staticmethod
    def map_int():
        return [int(x) for x in Scanner.string().split()]

    @staticmethod
    def string_list(n):
        return [input() for i in range(n)]

    @staticmethod
    def int_list_list(n):
        return [Scanner.map_int() for i in range(n)]

    @staticmethod
    def int_cols_list(n):
        return [int(input()) for i in range(n)]


class Math():
    @staticmethod
    def gcd(a, b):
        if b == 0:
            return a
        return Math.gcd(b, a % b)

    @staticmethod
    def lcm(a, b):
        return (a * b) // Math.gcd(a, b)

    @staticmethod
    def roundUp(a, b):
        return -(-a // b)

    @staticmethod
    def toUpperMultiple(a, x):
        return Math.roundUp(a, x) * x

    @staticmethod
    def toLowerMultiple(a, x):
        return (a // x) * x

    @staticmethod
    def nearPow2(n):
        if n <= 0:
            return 0
        if n & (n - 1) == 0:
            return n
        ret = 1
        while(n > 0):
            ret <<= 1
            n >>= 1
        return ret

    @staticmethod
    def sign(n):
        if n == 0:
            return 0
        if n < 0:
            return -1
        return 1

    @staticmethod
    def isPrime(n):
        if n < 2:
            return False
        if n == 2:
            return True
        if n % 2 == 0:
            return False
        d = int(n ** 0.5) + 1
        for i in range(3, d + 1, 2):
            if n % i == 0:
                return False
        return True


MOD = int(1e09) + 7
INF = int(1e15)

T_MAX = int(1e05) + 1


def main():
    # sys.stdin = open("sample.txt")
    N, K = Scanner.map_int()
    T = defaultdict(list)
    for _ in range(N):
        t, d = Scanner.map_int()
        t -= 1
        T[t].append(d)
    for v in T.values():
        v.sort(reverse=True)
    M = []
    for v in T.values():
        M.append((v[0], 1))
        for i in range(1, len(v)):
            M.append((v[i], 0))
    M.sort(reverse=True)
    q = []
    x = 0
    y = 0
    for i in range(K):
        x += M[i][0]
        if M[i][1] == 0:
            heapq.heappush(q, M[i][0])
        y += M[i][1]
    ans = x + y * y
    for i in range(K, N):
        if M[i][1] == 1:
            x -= q.pop()
            x += M[i][0]
            y += 1
            ans = max(ans, x + y * y)
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Runtime Error

Solution code:
Num,K=map(int,input().split())
L=[[]for i in range(Num+1)]
for i in range(Num):
    t,d=map(int,input().split())
    L[t].append(d)
for i in range(1,Num+1):
    L[i].sort(reverse=True)

M=[]
m=[]
for i in range(1,Num+1):
    if len(L[i])>=1:
        M.append(L[i][0])
        if len(L[i])>=2:
            for j in range(1,len(L[i])):
                m.append(L[i][j])
M.sort(reverse=True)
m.sort(reverse=True)
MM=[0]
for i in range(len(M)):
    MM.append(MM[-1]+M[i])
mm=[0]
for i in range(len(m)):
    mm.append(mm[-1]+m[i])
#print(MM)
#print(mm)
ANS=0

def cmb(n, r, mod):
    if ( r<0 or r>n ):
        return 0
    r = min(r, n-r)
    return g1[n] * g2[r] * g2[n-r] % mod

mod = 10**9+7 #出力の制限
N = 10**5+10
g1 = [1, 1] # 元テーブル
g2 = [1, 1] #逆元テーブル
inverse = [0, 1] #逆元テーブル計算用テーブル

for i in range( 2, N + 1 ):
    g1.append( ( g1[-1] * i ) % mod )
    inverse.append( ( -inverse[mod % i] * (mod//i) ) % mod )
    g2.append( (g2[-1] * inverse[-1]) % mod )

for i in range(len(M)+1):
    ANS=max(i*i+MM[i]+mm[min(len(mm)-1,K-i)],ANS)
    #print(ANS,i)
print(ANS)

Submission status: Runtime Error

Solution code:
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N+1)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    dt[i][1]-=1
    ans+=dt[i][0]
    if gr[dt[i][1]]==[]:
        x+=1
    gr[dt[i][1]].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    
inf=10**19
heapq.heappush(single,[inf,N])#絶対に選ばないものを入れておく
heapq.heappush(multi,[inf,N])
    
    
    
for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:#種類ボーナスも含めて小さい方を取り除く
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        tt=b[1]
        num_i[tt]+=1
        if num_i[tt]==len(gr[tt])-1:
            heapq.heappush(single,gr[tt][-1])
        else:
            heapq.heappush(multi,gr[tt][num_i[tt]])
                                       
print(ans)a

Submission status: Runtime Error

Solution code:
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    ans+=dt[i][0]
    if gr[dt[i][1]-1]==[]:
        x+=1
    gr[dt[i][1]-1].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    

for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        num_i[b[1]-1]+=1
        if num_i[b[1]-1]==len(gr[b[1]-1])-1:
            heapq.heappush(single,gr[b[1]-1][-1])
        else:
            heapq.heappush(multi,gr[b[1]-1][num_i[b[1]-1]])
                                       
print(ans)

Submission status: Runtime Error

Solution code:
def examA():
    L = LI(); L.sort()
    ans = L[0]*L[1]//2
    print(ans)
    return

def examB():
    def fun(s):
        if s%2==0:
            return s//2
        else:
            return s*3+1
    S = I()
    d = defaultdict(bool)
    d[S] = True; cur = 1
    while(True):
        cur +=1
        S = fun(S)
        if d[S]:
            break
        d[S] = True
    print(cur)
    return

def examC():
    N = I()
    H = LI()
    cur = 0
    ans =0
    for i in range(N):
        if cur<H[i]:
            ans += H[i]-cur
        cur = H[i]
    print(ans)
    return

def examD():
    N, K = LI()
    TS = [LI() for _ in range(N)]
    choice = []
    keep = []
    d = defaultdict(int)
    TS.sort(key=lambda x:x[1],reverse=True)
    for i,(t,s) in enumerate(TS):
        if d[t]==0:
            if i < K:
                d[t] = -1
                choice.append(s)
            else:
                d[t] = s
        else:
            if i<K:
                keep.append(s)
    keep.sort(reverse=True)
    D = sorted(d.values(),reverse=True)
#    print(choice); print(keep); print(D)
    cur = len(choice)
    now = cur**2 +sum(choice)+sum(keep)
    ans = copy.deepcopy(now)
    for i,k in enumerate(keep):
        now += ((cur+i)*2+1) + (D[i]-k)
        ans = max(ans,now)
#        print(now)
    print(ans)
    return

import sys,copy,bisect,itertools,heapq,math
from heapq import heappop,heappush,heapify
from collections import Counter,defaultdict,deque
def I(): return int(sys.stdin.readline())
def LI(): return list(map(int,sys.stdin.readline().split()))
def LSI(): return list(map(str,sys.stdin.readline().split()))
def LS(): return sys.stdin.readline().split()
def SI(): return sys.stdin.readline().strip()
global mod,mod2,inf,alphabet,ALPHABET, _ep
mod = 10**9 + 7
mod2 = 998244353
inf = 10**18
_ep = 1e-9
alphabet = [chr(ord('a') + i) for i in range(26)]
ALPHABET = [chr(ord('A') + i) for i in range(26)]

if __name__ == '__main__':
    examD()