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: Wrong Answer

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

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

print(answer)

Submission status: Wrong Answer

Solution code:

# -*- coding: utf-8 -*-
import bisect

N,K=map(int, input().split())
ts=[]
ds=[]
et={}


for i in range (0,K):
    t,d=map(int, input().split())
    
    ds.append(d)
    totmin='inf'

    #ds.append(d)
    bisect.insort(et.setdefault(t,[]),d)




for i in range (K,N):
    t,d=map(int, input().split())
    if t in et:
        dmin=et[t][0]
        if d>dmin:
            del et[t][0]
            bisect.insort(et[t],d)

            

    else:
        score=0
        place=-1
        l=len(et)
        for j in et.keys():
            dmin=et[j][0]

            if len(et[j])==1:
                if score<d-dmin:
                    place=j
                    score=d-dmin


            else:
                if score< d+2*l+1-dmin:
                    place=j
                    ddmin=dmin
                    score=d+2*l+1-dmin
            
        if place>=0:
            if len(et[place])==1:
                et.pop(place)
                et[t]=[d]

            else:
                et[t]=[d]
                del(et[place][0])


ans=len(et)**2
for i in et.keys():
    ans+=sum(et[i])

print(ans)

Submission status: Wrong Answer

Solution code:

# -*- coding: utf-8 -*-

N,K=map(int, input().split())
ts=[]
ds=[]
et={}


for i in range (0,K):
    t,d=map(int, input().split())
    
    ds.append(d)
    totmin='inf'

    #ds.append(d)
    et.setdefault(t,[]).append(d)




for i in range (K,N):
    t,d=map(int, input().split())
    if t in et:
        dmin=min(et[t])
        if d>dmin:
            et[t].remove(dmin)
            et[t].append(d)

            

    else:
        score=0
        place=-1
        l=len(et)
        for j in et.keys():
            dmin=min(et[j])

            if len(et[j])==1:
                if score<d-dmin:
                    place=j


            else:
                if score< d+2*l+1-dmin:
                    place=j
                    ddmin=dmin
            
        if place>=0:
            if len(et[place])==1:
                et.pop(place)
                et[t]=[d]

            else:
                et[t]=[d]
                et[place].remove(ddmin)


ans=len(et)**2
for i in et.keys():
    ans+=sum(et[i])

print(ans)

Submission status: Wrong Answer

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

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

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

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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

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

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

    

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    other_sushi += value

most_tasty.sort(reverse=True)
other_sushi.sort(reverse=True)
for i in range(1, k + 1):
    if i >= len(most_tasty):
        break
    point = sum(most_tasty[:i])
    point += sum(other_sushi[:(k - i)])
    point += i ** 2
    ans = max(ans, point)

print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    for v in value:
        other_sushi.append(v)

most_tasty.sort(reverse=True)
other_sushi.sort(reverse=True)
for i in range(1, k + 1):
    num = min(i, len(most_tasty))
    point = sum(most_tasty[:num])
    point += sum(other_sushi[:(k - num)])
    point += i ** 2
    ans = max(ans, point)

print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict

n, k = map(int, input().split())
ans = 0
sushi = defaultdict(list)

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

most_tasty = []
other_sushi = []

for key, value in sushi.items():
    value.sort(reverse=True)
    most_tasty.append(value[0])
    value.pop(0)
    for v in value:
        other_sushi.append(v)

most_tasty.sort(reverse=True)
for i in range(1, k + 1):
    num = min(i, len(most_tasty))
    point = sum(most_tasty[:num])
    point += sum(other_sushi[:(k - num)])
    point += i ** 2
    ans = max(ans, point)

print(ans)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

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:
    if c[td[ptr][0]] >= 1:
      ptr += 1
      continue
    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: Wrong Answer

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

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

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

point=f+len(V)**2

ans=point

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

Submission status: Wrong Answer

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

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

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

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

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

Submission status: Wrong Answer

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

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

eat = [[] for _ in range(N)]
types = set()
s = x = 0

# take the first K elements
for i in range(K):
    d, t = sushi[i]
    s += d
    if t not in types:
        types.add(t)
        x += 1
    eat[t].append(d)

# sushis that should be removed next
cand = []
for t in range(N):
    # if the type has more than one sushi
    if len(eat[t]) > 1:
        heappush(cand, (eat[t][::-1], t))

ans = s + x*x
for i in range(K, N):
    d, t = sushi[i]
    # if t is already eaten, no point eating worse one
    if t in types:
        continue
    # remove the worst sushi
    if not cand:
        break
    x += 1
    arr, tr = heappop(cand)
    dr = arr[0]
    s += d - dr
    ans = max(ans, s + x*x)
    eat[tr].pop()
    if len(eat[tr]) > 1:
        heappush(cand, (eat[tr][::-1], tr))
print(ans)

Submission status: Wrong Answer

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]))
if candel:
    qd = heapq.heappop(candel)
    while K+nex < N:
        if td[K+nex][0] not in kind:
            kind[td[K+nex][0]] = 1
            kinds += 1
            ans = ans - qd + td[K+nex][1]
            ans1 = max(ans1,ans+kinds**2)
            nex += 1
            if candel:
                qd = heapq.heappop(candel)
            else:
                break
        else:
            nex += 1
print(ans1)

Submission status: Wrong Answer

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

N, K = map(int, input().split())
sushi = [None] * N
for i in range(N):
    t, d = map(int, input().split())
    sushi[i] = (d, t-1)
sushi.sort(reverse=True)
sushi_used = [[] for _ in range(N)]
types = set()
s = 0
x = 0
for i in range(K):
    d, t = sushi[i]
    s += d
    if t not in types:
        types.add(t)
        x += 1
    sushi_used[t].append(d)
cand = []
for t in range(N):
    if len(sushi_used[t]) > 1:
        heappush(cand, (sushi_used[t][-1], t))
ans = s + x*x
for i in range(K, N):
    d, t = sushi[i]
    if t in types:
        continue
    x += 1
    if not cand:
        break
    dr, tr = heappop(cand)
    s += d - dr
    ans = max(ans, s + x*x)
    sushi_used[tr].pop()
    if len(sushi_used[tr]) > 1:
        heappush(cand, (sushi_used[tr][-1], tr))
    sushi_used[t].append(d)
    if len(sushi_used[t]) > 1:
        heappush(cand, (sushi_used[t][-1], t))
print(ans)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]

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

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

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

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

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

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

print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]

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

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

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

type_cnt = len(types)

ans = d_sm + type_cnt ** 2

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

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

print(ans)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

Solution code:
import sys
import heapq as h

input=sys.stdin.readline

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

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

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

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

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

Submission status: Wrong Answer

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

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

score, kind = 0, 0

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

ans = score + kind*kind

for t, d in l[K:]:
    if not t in used:
        heappush(q2, (-d, t))

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

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

print(ans)

Submission status: Wrong Answer

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

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

score = 0
for t, d in l[:K]:
    score += d
    used.add(t)
    heappush(q1[t], d)

kind = len(used)
ans = score + kind*kind
q2, q3 = [], []

for t, d in l[K:]:
    heappush(q2, (-d, t))

for t in q1.keys():
    c = len(q1[t])-1
    for _ in range(c):
        d = heappop(q1[t])
        heappush(q3, (d, t))

while q2 and q3:
    while q2:
        d1, t1 = heappop(q2)
        d1 = -d1
        if not t1 in used:
            break
    if not t1 is used:
        used.add(t1)
        d2, t2 = heappop(q3)
        score -= d2
        score += d1
        kind += 1
        ans = max(ans, score + kind*kind)

print(ans)

Submission status: Wrong Answer

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

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

score = 0
for t, d in l[:K]:
    score += d
    used.add(t)
    heappush(q1[t], d)

kind = len(used)
ans = score + kind*kind
q2, q3 = [], []

for t, d in l[K:]:
    heappush(q2, (-d, t))

for t in q1.keys():
    c = len(q1[t])-1
    for _ in range(c):
        d = heappop(q1[t])
        heappush(q3, (d, t))

while q2 and q3:
    while q2:
        d1, t1 = heappop(q2)
        d1 = -d1
        if not t1 in used:
            break
    if not t1 is used:
        d2, t2 = heappop(q3)
        score -= d2
        score += d1
        kind += 1
        ans = max(ans, score + kind*kind)

print(ans)

Submission status: Wrong Answer

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

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

score = 0
for t, d in l[:K]:
    score += d
    used.add(t)
    heappush(q1[t], d)

kind = len(used)
ans = score + kind*kind
q2, q3 = [], []

for t, d in l[K:]:
    heappush(q2, (-d, t))

for t in q1.keys():
    c = len(q1[t])-1
    for _ in range(c):
        d = heappop(q1[t])
        heappush(q3, (d, t))

while q2 and q3:
    while q2:
        d1, t1 = heappop(q2)
        d1 = -d1
        if not t1 in used:
            used.add(t1)
            break
    if not t1 is used:
        d2, t2 = heappop(q3)
        score -= d2
        score += d1
        kind += 1
        ans = max(ans, score + kind*kind)

print(ans)

Submission status: Wrong Answer

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

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

score = 0
for t, d in l[:K]:
    score += d
    heappush(q1[t], d)

used = set(q1.keys())
kind = len(used)

q2, q3 = [], []
for t, d in l[K:]:
    heappush(q2, (-d, t))

for t in q1.keys():
    c = len(q1[t])-1
    for _ in range(c):
        d = heappop(q1[t])
        heappush(q3, (d, t))

ans = score + kind*kind
while q2 and q3:
    while q2:
        d1, t1 = heappop(q2)
        d1 = -d1
        if not t1 in used:
            used.add(t1)
            break
    if not t1 is used:
        d2, t2 = heappop(q3)
        score -= d2
        score += d1
        kind += 1
        ans = max(ans, score + kind*kind)

print(ans)

Submission status: Wrong Answer

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))
    print(su+c*c)
else:
    ans = su+c*c
    
    while not_used:
        a, _ = heappop(now_used)
        su -= a
        a, _ = heappop(not_used)
        su -= a
        c -= 1
        ans = max(ans, su+c*c)
    print(ans)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

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

Submission status: Wrong Answer

Solution code:
import heapq

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

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

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

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

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

print(max(ans))

Submission status: Wrong Answer

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

Submission status: Wrong Answer

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

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

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

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

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

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

various = min(K, len(list1))
selected = []
for i in range(various):
    selected.append(list1.pop(0))
if K > various:
    for i in range(K - various):
        selected.append(list0.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: Wrong Answer

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 = min(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: Wrong Answer

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

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

base = []; other = []
val = set()
yum = 0
for t, d in td[:k]:
    val.add(t)
    yum += d
    if t not in val:
        base[t] = d
    else:
        other.append(d)
other.sort(reverse=True)

kouho = []
val_tmp = val.copy()
for t, d in td[k:]:
    if t not in val_tmp:
        kouho.append(d)
        val_tmp.add(t)
kouho.sort()

v = len(val)
ans = yum + v**2

while len(other) > 0 and len(kouho) > 0:
    yum -= other.pop()
    yum += kouho.pop()
    v += 1
    ans = max(ans, yum+(v**2))
print(ans)

Submission status: Wrong Answer

Solution code:

import heapq

n,k,*td=map(int,open(0).read().split())
td = [[t,d] for t,d in zip(td[::2],td[1::2])]
td = sorted(td,key = lambda x:x[1],reverse=True)

score = [0] * k
taste_sum = 0
extra_taste = []
eat_num = 0
eat_kind = set()

for t,d in td:
    if t in eat_kind:
        if eat_num == k:
            continue
        else:
            heapq.heappush(extra_taste, d)
            eat_num += 1
    else:
        eat_kind.add(t)
        taste_sum += d
        eat_num += 1
        if eat_num > k:
            eat_num -= 1
            extra_taste = extra_taste[1:]
    
    if eat_num == k:
        score[len(eat_kind)-1] = len(eat_kind)**2 + taste_sum + sum(extra_taste)
            
    if len(eat_kind) == k:
        break
print(max(score))

Submission status: Wrong Answer

Solution code:


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

from collections import defaultdict
import heapq

for n in range(N):
  t, d = map(int, input().split())
  sushi.append((d, t))
  kinds.append(d)

sushi = sorted(sushi, reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
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 and sushiq:
    d_min, t_min = heapq.heappop(sushiq)

  # popするものがなくなったら終了
  if not sushiq:
    break

  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: Wrong Answer

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

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

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)
kind_all = len(set(kinds))

cnt = 0
sushiq = []

for i in range(K):
  d, t = sushi[i]
  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:
    break

  d_min, t_min = heappop(sushiq)
  while sushidict[t_min] == 1 and sushiq:
    d_min, t_min = heappop(sushiq)

  if not sushiq:
    break

  sushidict[t_min] -= 1
  sushidict[t] += 1
  n_kind += 1
  cnt = cnt - d_min + d

  res.append(cnt+n_kind ** 2)
  if n_kind == kind_all:
    break

print(max(res))

Submission status: Wrong Answer

Solution code:
import sys
from collections import *
import heapq
import math
import bisect
from itertools import permutations,accumulate,combinations,product
from fractions import gcd
def input():
    return sys.stdin.readline()[:-1]
mod=10**9+7

n,k=map(int,input().split())
td=[list(map(int,input().split())) for i in range(n)]
td.sort(key=lambda x: -x[1])

dic={}
minlst=[]
maxlst=[]
ans=0
f=0
for i in range(n):
    t,d=td[i]
    if i<k:
        f+=d
        if not t in dic:
            dic[t]=d
        else:
            minlst.append(d)
    else:
        if not t in dic:
            maxlst.append(-d)
        else:
            pass

# print(td)
# print(minlst)
# print(maxlst)
# print(f)
# print(len(dic))
ln=len(dic)
ans=max(ans,f+pow(len(dic),2))
heapq.heapify(minlst)
heapq.heapify(maxlst)
while minlst and maxlst:
    f-=heapq.heappop(minlst)
    f-=heapq.heappop(maxlst)
    ln+=1
    ans=max(ans,f+pow(ln,2))
print(ans)