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:
import sys
input = sys.stdin.buffer.readline

from heapq import heappush, heappop

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

cnt_t = [0] * (N + 1)
take = [0] * N
TD.sort(key = lambda x: - x[1])

h = []

tmp = 0
kind = 0
for i, (t, d) in enumerate(TD):
  if cnt_t[t] == 0:
    cnt_t[t] += 1
    take[i] = 1
    kind += 1
    tmp += d
    heappush(h, (d, t))
rest = K - sum(cnt_t)

for i, (t, d) in enumerate(TD):
  if rest == 0:
    break
  if take[i] == 0:
    take[i] = 1
    cnt_t[t] += 1
    tmp += d
    rest -= 1
    heappush(h, (d, t))

rest = []
for i, (t, d) in enumerate(TD):
  if take[i] == 0:
    rest.append((d, t))
    
#print(h)
#print(rest)
answer = tmp + kind ** 2
for d_in, t_in in rest:
  d_out, t_out = heappop(h)
  tmp += d_in - d_out
  cnt_t[t_out] -= 1
  if cnt_t[t_out] == 0:
    kind -= 1
  cnt_t[t_in] += 1
  if cnt_t[t_in] == 1:
    kind += 1
  answer = max(answer, tmp + kind ** 2)
  heappush(h, (d_in, t_in))
  
print(answer)

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
ans = [0 for i in range(N+1)]
food = {}
food1 = []
for i in range(1,N+1):
    food[i] = 0
for i in range(N):
    t,d = map(int,input().split())
    food1.append((d,t))
food1.sort()
syurui = 0
sum = 0
x = []
for i in range(1,N+1):
    d,t = food1[-i]
    sum += d
    if food[t] == 0:
        food[t] = 1
        syurui += 1
        if i > K:
            y = x.pop(0)
            sum -= y
    else:
        x.append(d)
        x.sort()
    if syurui == K:
        ans[K] = max(sum + K**2,ans[K])
        break
    else:
        ans[syurui] = max(sum + syurui**2,ans[syurui])
print(max(ans))

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
food = {}
for i in range(1,N+1):
    food[i] = 0
food1 = []
for i in range(N):
    t,d = map(int,input().split())
    food[t] = max(food[t],d)
    food1.append(d)
food2 = []
for i in range(1,N+1):
    x = food[i]
    if x == 0:
        continue
    food1.remove(x)
    food2.append(-x)
food2.sort()
j = 1
for x in food2:
    x = -x
    x += j
    food1.append(x)
    j += 2
food1.sort()
ans = 0
for i in range(K):
    ans += food1.pop()
print(ans)

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
food = [[] for i in range(N+1)]
food1 = []
for i in range(N):
    t,d = map(int,input().split())
    food[t].append(d)
    food[t].sort()
    food1.append(d)
food2 = []
for i in range(1,N+1):
    if len(food[i]) == 0:
        continue
    x = food[i][-1]
    food1.remove(x)
    food2.append((-x,i))
food2.sort()
j = 1
for x,i in food2:
    x = -x
    x += j
    food1.append(x)
    j += 2
food1.sort()
ans = 0
for i in range(K):
    ans += food1.pop()
print(ans)

Submission status: Wrong Answer

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 = now
    for i,k in enumerate(keep):
#        print(now,((cur+i)*2+1))
        if D[i]<0:
            break
        now += ((cur+i)*2+1) + (D[i]-k)
        ans = max(ans,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()

Submission status: Wrong Answer

Solution code:
import sys
from heapq import heappush,heappop,heapify
sys.setrecursionlimit(10**9)
INF=10**18
MOD=10**9+7
def input():
    return sys.stdin.readline().rstrip()

def main():
    N,K=map(int,input().split())
    sushies=[]
    for i in range(N):
        t,d=map(int,input().split())
        sushies.append((d,t-1))
    sushies.sort()
    b=[False]*N
    overlappedsushi=[]
    ans=0
    n=0
    for i in range(K):
        sushi=sushies.pop()
        if b[sushi[1]]:
            heappush(overlappedsushi,sushi)
            ans+=sushi[0]
        else:
            b[sushi[1]]=True
            ans+=sushi[0]
            n+=1
    ans+=n**2
    tmp=ans
    while overlappedsushi:
        worstsushi=heappop(overlappedsushi)
        tmp-=worstsushi[0]
        while sushies:
            bestsushi=sushies.pop()
            if not b[bestsushi[1]]:
                break
        else:
            break
        n+=1
        tmp+=bestsushi[0]
        tmp+=n*2-1
        ans=max(ans,tmp)
    print(ans)

    

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import heapq

N,K = map(int, input().split())
h = []
max_t = 0
for _ in range(N):
    t, d = map(int, input().split())
    heapq.heappush(h, (-d, t))
    max_t = max(max_t, t)

L = [0]*max_t
num = 0
h2 = []
for _ in range(K):
    temp = heapq.heappop(h)
    L[temp[1]-1] += 1
    num -= temp[0]
    heapq.heappush(h2, (-temp[0], temp[1]))
    
cnt = 0
for i in range(max_t):
    if L[i] > 0:
        cnt += 1

ans = num + cnt**2
while len(h) > 0 and len(h2) > 0:
    temp = heapq.heappop(h)
    if L[temp[1]-1] > 0:
        continue
    if L[h2[0][1]-1] > 1:
        temp2 = heapq.heappop(h2)
    else:
        continue
        
    num -= temp2[0] + temp[0]
    L[temp2[1]-1] -= 1
    L[temp[1]-1] += 1
    cnt += 1
    ans = max(ans, num+cnt**2)

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
from collections import defaultdict as dd
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)
for i in range(N):
  if table[s[i][0]] == (-1, -1): table[s[i][0]] = (i, s[i][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, M + 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: Wrong Answer

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 and len(sum_o) > 0:
            tmp += sum_o[min(K-(i+1),len(sum_o))-1]
        ans = max(ans, tmp)

    print(ans)

    
    




if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-

import sys
from heapq import heapify, heappush, heappop
from collections import defaultdict

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10 ** 9)
INF = 10 ** 18
MOD = 10 ** 9 + 7

N, K = MAP()
TD = []
tmp = [[] for i in range(N)]
for i in range(N):
    t, d = MAP()
    t -= 1
    tmp[t].append(d)

for li in tmp:
    li.sort(reverse=True)
cur = [0, defaultdict(int)]
heapa = []
heapb = []
for t, li in enumerate(tmp):
    if li:
        cur[0] += li[0]
        cur[1][t] = 1
        heapa.append((li[0], t))
        for d in li[1:]:
            TD.append((t, d))
            heapb.append((-d, t))

heapa.sort(reverse=1)
for i in range(K, len(heapa)):
    d, t = heapa[i]
    cur[0] -= d
    cur[1][t] -= 1
    if cur[1][t] == 0:
        del cur[1][t]
    heapb.append((-d, t))
heapa = heapa[:K]

heapify(heapa)
heapify(heapb)

M = len(cur[1])
ans = cur[0] + len(cur[1])**2
for i in range(M-1):
    d, t = heappop(heapa)
    cur[0] -= d
    cur[1][t] -= 1
    if cur[1][t] == 0:
        del cur[1][t]
    if not heapb:
        break
    d, t = heappop(heapb)
    d = -d
    cur[0] += d
    cur[1][t] += 1
    ans = max(ans, cur[0] + len(cur[1])**2)
print(ans)

Submission status: Wrong Answer

Solution code:
from bisect import bisect_left

n,k = map(int,input().split())
td = [list(map(int,input().split())) for i in range(n)]

mt = max(td)[0]
# print(mt)
one = [0]*(mt+1)
zero = []
for t,d in td:
    if one[t] < d:
        zero.append(one[t])
        one[t] = d
    else:
        zero.append(d)

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

idx = bisect_left(one,1)
cur = sum(one[idx:idx+k+1])
maisuu = len(one[idx:idx+k+1])
syurui = maisuu
ans = cur+maisuu**2
# print(cur,maisuu,ans)

for z in zero:
    if maisuu != k:
        cur += 1
        maisuu += 1
        continue
    o = one[idx]
    idx += 1
    cur -= o
    cur += z
    syurui -= 1
    # print(ans,cur,syurui)
    if ans > cur+syurui**2:
        break
    ans = cur+syurui**2
print(ans)

Submission status: Wrong Answer

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[0],x[1]),reverse=1)

new_td = []
before = 0
for t,d in td:
    if t != before:
        new_td.append([d,1])
        before = t
    else:
        new_td.append([d,0])

new_td.sort(reverse=1)

ans = 0
maisuu = 0
cur = 0
zero = []
for val,pea in new_td:
    if pea == 0:
        zero = [val]+zero
    if maisuu != k:
        cur += val
        maisuu += 1
        ans = cur+(maisuu-len(zero))**2
    else:
        if not zero:
            break
        cur += val-zero.pop()
        ans = max(cur+(k-len(zero))**2,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
def solve(n, k, t, d):
    que = []
    used = set()
    base_score = 0
    for ti, di in sorted(zip(t, d), key=lambda _:-_[1])[:k]:
        if ti in used:
            que.append(di)
        else:
            used.add(ti)
        base_score += di
    num = len(used)
    score = base_score + num**2
    best = score

    M = {_:0 for _ in t}
    for ti, di in zip(t, d):
        if not ti in used:
            M[ti] = max(M[ti], di)

    que = sorted(que)
    for i, p in enumerate(sorted(M.values())[::-1][:len(que)]):
        base_score += p - que[i]
        num += 1
        best = max(best, base_score + num**2)
    return best

n, k = map(int, input().split())
t = [0] * n
d = [0] * n
for i in range(n):
    t[i], d[i] = map(int, input().split())
print(solve(n, k, t, d))

Submission status: Wrong Answer

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
while 1:
    try:
        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)
    except IndexError:
        break

print(max(ans_list))

Submission status: Wrong Answer

Solution code:
import sys
from operator import itemgetter
read = sys.stdin.read

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

stack = []
taste = 0
neta_group = set()
for t, d in td[:K]:
    taste += d
    if t in neta_group:
        stack.append(d)
    else:
        neta_group.add(t)

variety = len(neta_group)
points = taste + variety ** 2
candidates = [points]

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

rest = list(reversed(td[K::]))

while rest and stack:
    print(points, variety)
    t1, d1 = rest.pop()
    d2 = stack.pop()
    points = points - d2 + d1 - variety ** 2 + (variety + 1) ** 2

    candidates.append(points)
    variety += 1

print(max(candidates))

Submission status: Wrong Answer

Solution code:
import os, sys, re, math

(N, K) = [int(n) for n in input().split()]
TD = []
counts = [0] * (N + 1)
for i in range(N):
    td = [int(n) for n in input().split()]
    counts[td[0]] += 1
    TD.append(td)

by_delicious = sorted(TD, key=lambda x:x[1], reverse=True)

points = 0
ingredients = set()
for i in range(K):
    points += TD[i][1]
    ingredients.add(TD[i][0])

d_point = points + len(ingredients) ** 2

for i in range(len(TD)):
    TD[i] = TD[i] + [counts[TD[i][0]]]

by_rare = sorted(TD, key=lambda x:x[2])

points = 0
ingredients = set()
for i in range(K):
    points += TD[i][1]
    ingredients.add(TD[i][0])

r_point = points + len(ingredients) ** 2

print (max(d_point, r_point))

Submission status: Wrong Answer

Solution code:
import os, sys, re, math

(N, K) = [int(n) for n in input().split()]
TD = []
for i in range(N):
    TD.append([int(n) for n in input().split()])

by_delicious = sorted(TD, key=lambda x:x[1], reverse=True)

points = 0
ingredients = set()
for i in range(K):
    points += TD[i][1]
    ingredients.add(TD[i][0])

print (points + len(ingredients) ** 2)

Submission status: Wrong Answer

Solution code:
from heapq import heappop, heappush

n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x: (-x[1], x[0]))

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    type, taste = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        heappush(same, taste)

rem = []
typesRem = set()
for i in range(nTake, n):
    type, taste = a[i]
    if type not in types and type not in typesRem:
        heappush(rem, -taste)

ret = totalTaste + len(types) ** 2
curTotal = totalTaste
count = 1
while same and rem:
    worst = heappop(same)
    best = heappop(rem)
    curTotal = curTotal - worst - best
    ret = max(ret, curTotal + (len(types) + count) ** 2)
    count += 1
print(ret)

Submission status: Wrong Answer

Solution code:
from heapq import heappop, heappush

n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split()))[::-1] for _ in range(n)], key=lambda x: (-x[0], x[1]))

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    taste, type = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        heappush(same, taste)

left = []
for i in range(nTake, n):
    taste, type = a[i]
    if type not in types:
        heappush(left, -taste)

ret = totalTaste + len(types) ** 2
curTotal = totalTaste
count = 1
while same and left:
    worst = heappop(same)
    best = heappop(left)
    curTotal = curTotal - worst - best
    ret = max(ret, curTotal + (len(types) + count) ** 2)
    count += 1
print(ret)

Submission status: Wrong Answer

Solution code:
from collections import*
n,k=map(int,input().split())
sushi=sorted([tuple(map(int,input().split()))for i in range(n)],key=lambda x:x[1],reverse=1)
first=defaultdict(list)
second=defaultdict(list)
menu=set()

for k,v in sushi[:k]:
    menu.add(k)
    first[k].append(v)
second=sushi[k:][::-1]
#print(sushi,first,second)

score=sum(map(sum,first.values()))
change=[]
for k,v in first.items():
    for i in v[1:]:
        change.append((k,i))
change.sort(key=lambda x:x[1],reverse=1)


ans=[score+len(first)**2]
while change:
    point,s=change.pop()
    score-=s
    while 1:
        if second:
            if second[-1][0] in menu:
                pp,ss=second.pop()
                ans.append(score+ss+len(menu)**2)
            else:
                break
        else:
            break
    if second:
        pp,ss=second.pop()
    else:
        break
    menu.add(pp)
    score+=ss
    ans.append(score+len(menu)**2)
print(max(ans))

Submission status: Wrong Answer

Solution code:
from collections import*
n,k=map(int,input().split())
sushi=sorted([tuple(map(int,input().split()))for i in range(n)],key=lambda x:x[1],reverse=1)
first=defaultdict(list)
second=defaultdict(list)
menu=set()

for k,v in sushi[:k]:
    menu.add(k)
    first[k].append(v)
second=sushi[k:][::-1]
#print(sushi,first,second)

score=sum(map(sum,first.values()))
change=[]
for k,v in first.items():
    for i in v[1:]:
        change.append((k,i))
change.sort(key=lambda x:x[1],reverse=1)


ans=[score+len(first)**2]
while change and second:
    point,s=change.pop()
    score-=s
    try:
        while second[-1][0] in menu:
                pp,ss=second.pop()
        pp,ss=second.pop()
        menu.add(pp)
        score+=ss
        ans.append(score+len(menu)**2)
    except:
        break
print(max(ans))

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop

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

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

def solve() :
    done = [False] * (N + 1)
    k = min(K, len(set([td[i][0] for i in range(N)])))
    
    base = 0
    kind = 0
    count = 0
    for t, d in td :
        if done[t] :
            if d > 2 * k - 1 :
                base += d
                count += 1
                k -= 1
        else :
            base += d
            kind += 1
            count += 1
            done[t] = True
        
        if count == K :
            return base + kind * kind
            
print(solve())

Submission status: Wrong Answer

Solution code:
import sys
input = sys.stdin.readline
sys.setrecursionlimit(pow(10, 6))
from heapq import heappush, heappop

def main():
    n, k = map(int, input().split())
    td = [tuple(map(int, input().split())) for _ in range(n)]
    td = list(sorted(td, key=lambda x: -x[1]))
    dsum, kind = 0, 0
    num = [0 for _ in range(n)]
    h = []
    for t, d in td[:k]:
        dsum += d
        if num[t-1] > 0:
            heappush(h, (-d, t))
        if num[t-1] == 0:
            kind += 1
        num[t-1] += 1
    
    ans = dsum + kind * kind

    for t, d in td[k:]:
        if len(h) == 0:
            break
        if num[t-1] > 0:
            continue
        else:
            _d, _t = heappop(h)
            num[_t-1] -= 1
            num[t-1] += 1
            kind += 1
            dsum += _d+d
            if dsum + kind * kind > ans:
                ans = dsum + kind * kind
    
    print(ans)


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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


def main():
    n, k = map(int, input().split())
    td = []
    td_m = {}
    ts = set()
    for _ in range(n):
        t, d = map(int, input().split())
        td.append([t, d])
        td_m[t] = []
        heapq.heapify(td_m[t])
        ts.add(t)

    for t, d in td:
        heapq.heappush(td_m[t], d)

    m = {0: [], 1: []}
    for i in ts:
        if len(td_m[i]) == 1:
            m[1].append(td_m[i][0])
        else:
            for x in td_m[i][:-1]:
                m[0].append(x)

            m[1].append(td_m[i][-1])

    m[0].sort(reverse=True)
    m[1].sort(reverse=True)

    ans = []
    for i in range(1, len(m[1])):
        ans.append(sum(m[1][:i]) + sum(m[0][:(k-i)]) + i * i)
    ans.append(sum(m[1]) + len(m[1]) * len(m[1]))

    print(max(ans))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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


def main():
    n, k = map(int, input().split())
    td = []
    td_m = {}
    ts = set()
    for _ in range(n):
        t, d = map(int, input().split())
        td.append([t, d])
        td_m[t] = []
        heapq.heapify(td_m[t])
        ts.add(t)

    for t, d in td:
        heapq.heappush(td_m[t], d)

    m = {0: [], 1: []}
    for i in ts:
        if len(td_m[i]) == 1:
            m[1].append(td_m[i][0])
        else:
            for x in td_m[i][:-1]:
                m[0].append(x)

            m[1].append(td_m[i][-1])

    m[0].sort(reverse=True)
    m[1].sort(reverse=True)

    ans = []
    for i in range(1, len(m[1])):
        ans.append(sum(m[1][:i]) + sum(m[0][:(k-i)]) + i * i)
    ans.append(sum(m[1]) + len(m[1] * len(m[1])))

    print(max(ans))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
from heapq import heappop, heappush
 
N, K = map(int, input().split())
td = sorted([list(map(int, input().split()))
             for _ in range(N)], key=lambda x: x[1], reverse=True)

q = []
v = set()
count = 0
for t, d in td[:K]:
    count += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
count += len(v) ** 2
ans = count
for t, d in td[:K]:
    if t not in v and len(q) != 0:
        x = heappop(q)
        count += d + 2 * len(v) + 1 - x
        v.add(t)
        ans = max(ans, count)
        
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append([d,t])
data.sort(reverse=True)
flag=[0]*(n+1)
eat=[]
lst=[]
count=0
for u in data:
    d,t=u
    if flag[t]==0:
        eat.append(u)
        flag[t]=1
        count+=1
        if count==k:
            break
    else:
        lst.append(u)
else:
    lst.sort(reverse=True)
    sum=0
    for i in range(k-count):
        sum+=lst[i][0]
    for u in eat:
        sum+=u[0]
    sum+=count**2
    print(sum)
    print("y")
    sys.exit()
lst.sort(reverse=True)
for u in lst:
    d=u[0]
    if d>eat[count-1][0]+2*count-1:
        eat[count-1]=u
        count-=1
    else:
        break
sum=0
for u in eat:
    sum+=u[0]
sum+=count**2
print(sum)

Submission status: Wrong Answer

Solution code:
import sys
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append([d,t])
data.sort(reverse=True)
flag=[0]*(n+1)
eat=[]
lst=[]
count=0
for u in data:
    d,t=u
    if flag[t]==0:
        eat.append(u)
        flag[t]=1
        count+=1
        if count==k:
            break
    else:
        lst.append(u)
else:
    lst.sort(reverse=True)
    sum=0
    for i in range(k-count):
        sum+=lst[i][0]
    for u in eat:
        sum+=u[0]
    sum+=count**2
    print(sum)
    sys.exit()
lst.sort(reverse=True)
for u in lst:
    d=u[0]
    if d>eat[count-1][0]+2*count-1:
        eat[count-1]=u
        count-=1
    else:
        break
sum=0
for u in eat:
    sum+=u[0]
sum+=count**2
print(sum)

Submission status: Wrong Answer

Solution code:
import heapq

INF=10**12

N, K = map(int,input().split())
td = [[0,0] for i in range(N)]
s = [[] for i in range(N)]

for i in range(N):
    td[i][0], td[i][1] = map(int,input().split())
    td[i][0] -= 1
    s[td[i][0]].append(td[i][1])

for i in range(N):
    if s[i] == []:
        s[i].append(-INF)
    s[i].sort(key=lambda x:-x)

s.sort(key=lambda x:-x[0])

tmp = 0
q = []

for i in range(K):
    tmp += s[i][0]
    for j in range(1,len(s[i])):
        heapq.heappush(q,-s[i][j])
for i in range(K,N):
    for j in range(len(s[i])):
        heapq.heappush(q,-s[i][j])

res = tmp + K * K

for i in range(1,K)[::-1]:
    v = s[i][0]
    w = - q[0]
    if v < w:
        tmp = tmp + w - v
        heapq.heappop(q)
        heapq.heappush(q,v)
    res = max(res, tmp + i * i)        

print(res)

Submission status: Wrong Answer

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

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 ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

res = cur + K * K
for x in reversed(range(1, K)):
    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: Wrong Answer

Solution code:
from heapq import heappush, heappop, heapreplace

INF = float("inf")

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

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

for e in E:
    if e:
        e.sort(reverse=True)
    else:
        e.append(-INF)
E.sort(key=lambda x: -x[0])

cur = 0
Q = []
for ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

res = cur + K * K
for x in reversed(range(1, K)):
    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: Wrong Answer

Solution code:
from heapq import heappush, heappop, heapreplace

INF = float("inf")

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

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

for e in E:
    if e:
        e.sort(reverse=True)
    else:
        e.append(-INF)
E.sort(key=lambda x: -x[0])

cur = 0
Q = []
for ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

res = cur + K * K
for x in reversed(range(1, K)):
    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: Wrong Answer

Solution code:
from heapq import heappush, heapreplace

INF = float("inf")

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

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

for e in E:
    if e:
        e.sort(reverse=True)
    else:
        e.append(-INF)
E.sort(key=lambda x: -x[0])

cur = 0
Q = []
for ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

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

print(res)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappushpop

INF = 1e40

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

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

for e in E:
    if e:
        e.sort(reverse=True)
    else:
        e.append(-INF)
E.sort(key=lambda x: -x[0])

cur = 0
Q = []
for ei in E[:K]:
    cur += ei[0]
    for eij in ei[1:]:
        heappush(Q, -eij)

for ei in E[K:]:
    for eij in ei:
        heappush(Q, -eij)

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

print(res)