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 collections import defaultdict, deque, Counter
from heapq import heappush, heappop, heapify
import math
import bisect
import random
from itertools import permutations, accumulate, combinations, product
import sys
import string
from bisect import bisect_left, bisect_right
from math import factorial, ceil, floor
from operator import mul
from functools import reduce


INF = float('inf')
def LI(): return list(map(int, sys.stdin.readline().split()))
def I(): return int(sys.stdin.readline())
def LS(): return sys.stdin.readline().split()
def S(): return sys.stdin.readline().strip()
def IR(n): return [I() for i in range(n)]
def LIR(n): return [LI() for i in range(n)]
def SR(n): return [S() for i in range(n)]
def LSR(n): return [LS() for i in range(n)]
def SRL(n): return [list(S()) for i in range(n)]
def MSRL(n): return [[int(j) for j in list(S())] for i in range(n)]
mod = 1000000007


n, k = LI()
L = []
for kind, point in LIR(n):
    L += [(point, kind)]
initial_menu = sorted(L)[-k:]
ret = sum([i for i, j in initial_menu])
kind_cnt = Counter([i[1] for i in initial_menu])
now_kinds = len(kind_cnt)
ans = ret + now_kinds ** 2
que1 = deque(initial_menu)
que2 = deque(L[:k])
while que1:
    point1, kind1 = que1.popleft()
    if kind_cnt[kind1] > 1:
        kind_cnt[kind1] -= 1
        ret -= point1
        while que2:
            point2, kind2 = que2.pop()
            if not kind_cnt[kind2]:
                ret += point2
                now_kinds += 1
                ans = max(ans, ret + now_kinds ** 2)
                break
        else:
            break



print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict, deque, Counter
from heapq import heappush, heappop, heapify
import math
import bisect
import random
from itertools import permutations, accumulate, combinations, product
import sys
import string
from bisect import bisect_left, bisect_right
from math import factorial, ceil, floor
from operator import mul
from functools import reduce


INF = float('inf')
def LI(): return list(map(int, sys.stdin.readline().split()))
def I(): return int(sys.stdin.readline())
def LS(): return sys.stdin.readline().split()
def S(): return sys.stdin.readline().strip()
def IR(n): return [I() for i in range(n)]
def LIR(n): return [LI() for i in range(n)]
def SR(n): return [S() for i in range(n)]
def LSR(n): return [LS() for i in range(n)]
def SRL(n): return [list(S()) for i in range(n)]
def MSRL(n): return [[int(j) for j in list(S())] for i in range(n)]
mod = 1000000007


n, k = LI()
L = []
for kind, point in LIR(n):
    L += [(point, kind)]
initial_menu = sorted(L)[-k:]
ret = sum([i for i, j in initial_menu])
kind_cnt = Counter([i[1] for i in initial_menu])
now_kinds = len(kind_cnt)
ans = ret + now_kinds ** 2
que1 = deque(initial_menu)
que2 = deque(L[:k])
while que1:
    point1, kind1 = que1.popleft()
    if kind_cnt[kind1] > 1:
        ret -= point1
        while que2:
            point2, kind2 = que2.pop()
            if not kind_cnt[kind2]:
                kind_cnt[kind2] -= 1
                ret += point2
                now_kinds += 1
                ans = max(ans, ret + now_kinds ** 2)
                break
        else:
            break



print(ans)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict, deque, Counter
from heapq import heappush, heappop, heapify
import math
import bisect
import random
from itertools import permutations, accumulate, combinations, product
import sys
import string
from bisect import bisect_left, bisect_right
from math import factorial, ceil, floor
from operator import mul
from functools import reduce


INF = float('inf')
def LI(): return list(map(int, sys.stdin.readline().split()))
def I(): return int(sys.stdin.readline())
def LS(): return sys.stdin.readline().split()
def S(): return sys.stdin.readline().strip()
def IR(n): return [I() for i in range(n)]
def LIR(n): return [LI() for i in range(n)]
def SR(n): return [S() for i in range(n)]
def LSR(n): return [LS() for i in range(n)]
def SRL(n): return [list(S()) for i in range(n)]
def MSRL(n): return [[int(j) for j in list(S())] for i in range(n)]
mod = 1000000007


n, k = LI()
L = []
for kind, point in LIR(n):
    L += [(point, kind)]
initial_menu = sorted(L)[-k:]
ret = sum([i for i, j in initial_menu])
kind_cnt = Counter([i[1] for i in initial_menu])
now_kinds = len(kind_cnt)
ans = ret + now_kinds ** 2
que1 = deque(initial_menu)
que2 = deque(L[:k])
while que1:
    point1, kind1 = que1.popleft()
    if kind_cnt[kind1] > 1:
        ret -= point1
        while que2:
            point2, kind2 = que2.pop()
            if not kind_cnt[kind2]:
                ret += point2
                now_kinds += 1
                ans = max(ans, ret + now_kinds ** 2)
                break
        else:
            break



print(ans)

Submission status: Wrong Answer

Solution code:
#00:46
n,k = map(int,input().split())
raw = [[] for _ in range(n)]
for _ in range(n):
  t,d = map(int,input().split())
  t -= 1
  raw[t].append(d)
only = []
other = []
for i in range(n):
  if len(raw[i]) != 0:
    raw[i].sort(reverse = True)
    only.append(raw[i][0])
    other += raw[i][1:]
only.sort(reverse = True)
other.sort(reverse = True)
lk = min(len(only),k)
now = sum(only[:lk])
can = [now + k * k]
for j in range(k):
  if j < k-lk:
    continue
  try:
    now -= only[k-1-j]
    now += other[j]
    can.append(now + (k-1-j) ** 2)
  except:
    break
print(max(can))

Submission status: Wrong Answer

Solution code:
#00:46
n,k = map(int,input().split())
raw = [[] for _ in range(n)]
for _ in range(n):
  t,d = map(int,input().split())
  t -= 1
  raw[t].append(d)
only = []
other = []
for i in range(n):
  if len(raw[i]) != 0:
    raw[i].sort(reverse = True)
    only.append(raw[i][0])
    other += raw[i][1:]
only.sort(reverse = True)
only += [0 for _ in range(k)]
other.sort(reverse = True)
now = sum(only[:k])
can = [now + k * k]
for j in range(k):
  try:
    now -= only[k-1-j]
    now += other[j]
    can.append(now + (k-1-j) ** 2)
  except:
    break
print(max(can))

Submission status: Wrong Answer

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

kinds = []
ans = 0
for i in range(k):
  if not sushi[0][0] in kinds:
    ans += sushi[0][1]
    kinds.append(sushi[0][0])
    sushi.remove(sushi[0])
  else:
    for j in range(1, len(sushi)):
      if not sushi[j][0] in kinds:
        if sushi[j][1] + 2* len(kinds) + 1 >= sushi[0][1]:
          ans += sushi[j][1]
          kinds.append(sushi[j][0])
          sushi.remove(sushi[j])
        else:
          ans += sushi[0][1]
          sushi.remove(sushi[0])
        break
print(ans + (len(kinds) ** 2))

Submission status: Wrong Answer

Solution code:
from heapq import heappop, heappush, heapify
n, k = map(int, input().split())
lst = [list(map(int, input().split())) for i in range(n)]
lst = sorted(lst, reverse=True, key=lambda x:x[1])
q = []
v = set()
cnt = 0
for t, d in lst[:k]:
    cnt += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
cnt += len(v)**2
ans = cnt
for t, d in lst[k:]:
    if t not in v and len(q)!=0:
        x = heappop(q)
        cnt += d + 2*len(v) + 1 - x
        ans = max(ans, cnt)
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[0],x[1]),reverse=True)
x= []
y= []
y.append(td[0][1])
for i in range(1,N):
    if td[i][0] != td[i-1][0]:
        y.append(td[i][1])
    else:
        x.append(td[i][1])
x.sort(reverse=True)
y.sort(reverse=True)
ans = 0
for i in range(1,len(y)+1):
    ans = max(ans, sum(y[:i])+sum(x[:max(0,K-i)])+i*i)
print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(N)]
A.sort(key = lambda x: -x[1])
res = 0
def check(n):
  t = 0
  k = 0
  result = 0
  visited = [0] * (N + 1)
  for i in range(N):
    #print(result)
    if visited[A[i][0]] == 0:
      if t < n:
        t += 1
        k += 1
        result += A[i][1]
        visited[A[i][0]] = 1
    else:
      if K - k > n - t:
        k += 1
        result += A[i][1]
    if k > K:
      break
  return result
def bisectS(l, r):
  global res
  mid = (l + r) // 2
  temp = check(mid) + mid ** 2
  #print(l, r, mid, temp)
  if abs(l - r) < 2:
    return
  if temp >= res:
    res = temp
    bisectS(mid, r)
  else:
    bisectS(l, mid)
res = check(K // 2) + (K // 2) ** 2
bisectS(1, K + 1)
print(res)

Submission status: Wrong Answer

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

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
a=sorted([list(map(int,input().split()))for i in range(n)],key=lambda x:-x[1])
t=[0]*(n+1)
p,q=0,0
for i in a[:k]:
    p+=i[1]
    if t[i[0]]==0:
        q+=1
    t[i[0]]+=1
s=0
d=p+q*q
c=a[:k][::-1]
for i in a[k:]:
    if t[i[0]]==0:
        while s!=k:
            if t[c[s][0]]<2:
                s+=1
            else:
                break
        if s==k:
            break
        else:
            t[i[0]]+=1
            p-=c[s][1]-i[1]
            q+=1
            d=max(d,p+q*q)
print(d)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict as de
n,k=map(int,input().split())
a=sorted([list(map(int,input().split()))for i in range(n)],reverse=True)
t=de(int)
p,q=0,0
for i in a[:k]:
    p+=i[1]
    if t[i[0]]==0:
        q+=1
    t[i[0]]+=1
s=0
d=p+q*q
c=a[:k][::-1]
for i in a[k:]:
    if t[i[0]]==0:
        while s!=k:
            if t[a[s][0]]>0:
                s+=1
            else:
                break
        t[i[0]]+=1
        p-=a[s][1]-i[1]
        q+=1
        d=max(d,p+q*q)
print(d)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict as de
n,k=map(int,input().split())
a=sorted([list(map(int,input().split()))for i in range(n)],reverse=True)
t=de(int)
p,q=0,0
for i in a[:k]:
    p+=i[0]
    if t[i[1]]==0:
        q+=1
    t[i[1]]+=1
s=0
d=p+q*q
c=a[:k][::-1]
for i in a[k:]:
    if t[i[1]]==0:
        while s!=k:
            if t[a[s][1]]>0:
                s+=1
            else:
                break
        t[i[1]]+=1
        p-=a[s][0]-i[0]
        q+=1
        d=max(d,p+q*q)
print(d)

Submission status: Wrong Answer

Solution code:
import heapq
n,k = (int(x) for x in input().split())
arr = []
discard = []
type = set()
score = 0

for i in range(n):
  a,b = (int(x) for x in input().split())
  heapq.heappush(arr, (-b, a))
  
# 初期状態
for i in range(k):
  d, t = heapq.heappop(arr)
  score -= d
  if t not in type:
    type.add(t)
  else:
    heapq.heappush(discard, (-d, t))

bonus = len(type)
ans = score + bonus**2

while arr and discard:
  pd, pt = heapq.heappop(arr)  #追加候補
  if pt not in type:
    _pd, _pt = heapq.heappop(discard) #被っているものの中で美味しさが最も小さいものを取り出す
    type.add(_pt)
    score = score - _pd - pd
    bonus += 1
    tmp = score + bonus**2
    ans = max(ans, tmp)
    
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
n,k = (int(x) for x in input().split())
arr = []
discard = []
type = set()
score = 0

for i in range(n):
  a,b = (int(x) for x in input().split())
  heapq.heappush(arr, (-b, a))
  
# 初期状態
for i in range(k):
  d, t = heapq.heappop(arr)
  score -= d
  if t not in type:
    type.add(t)
  else:
    heapq.heappush(discard, (-d, t))

bonus = len(type)
ans = score + bonus**2

while arr and discard:
  pd, pt = heapq.heappop(arr)  #追加候補
  if pt not in type:
    _pd, _pt = heapq.heappop(discard) #被っているものの中で美味しさが最も小さいものを取り出す
    score = score - _pd - pd
    bonus += 1
    tmp = score + bonus**2
    ans = max(ans, tmp)
    
print(ans)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop
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)
muptiple = set()
hp = []
for t, d in TD:
    flg = 0
    if t not in muptiple:
        muptiple.add(t)
        flg = 1
    heappush(hp, (d*-1, flg))

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = heappop(hp)
    t = v[0]*-1
    vari += v[1]
    point += t
    heappush(ret, (v[1]*(10**9)+t, t, v[1]))

total = vari ** 2 + point

while len(hp) > 0:
    v = heappop(hp)
    if v[1] == 0:
        continue
    pre_v = heappop(ret)
    t = v[0]*-1
    heappush(ret, (v[1]*(10**10)+t, t, v[1]))
    point -= pre_v[1] - t
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop
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)
muptiple = set()
hp = []
for t, d in TD:
    flg = 0
    if t not in muptiple:
        muptiple.add(t)
        flg = 1
    heappush(hp, (d*-1, flg))

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = heappop(hp)
    t = v[0]*-1
    vari += v[1]
    point += t
    heappush(ret, (v[1]*(10**9)+t, t, v[1]))

total = vari ** 2 + point

while len(hp) > 0:
    v = heappop(hp)
    if v[1] == 0:
        continue
    pre_v = heappop(ret)
    t = v[0]*-1
    heappush(ret, (v[1]*(10**9)+t, t, v[1]))
    point -= pre_v[1] - t
    vari += 1
    tmp = vari ** 2 + point
    total = max(total, tmp)

print(total)

Submission status: Wrong Answer

Solution code:
def main():

    """ convenient functions
    # for i, a in enumerate(iterable)
    # q, mod = divmod(a, b)
    # divmod(x, y) returns the tuple (x//y, x%y)
    # Higher-order function: reduce(operator.mul, xyz_count, 1)
    # manage median(s) using two heapq https://atcoder.jp/contests/abc127/tasks/abc127_f
    """

    """convenient decorator
    # @functools.lru_cache():
    # to facilitate use of recursive function
        # ex:
        # from functools import lru_cache
        # import sys
        # sys.setrecursionlimit(10**9)
        # @lru_cache(maxsize=None)
        # def fib(n):
        #     if n < 2:
        #         return n
        #     return fib(n-1) + fib(n-2)
        # print(fib(1000))
    """
    
    # import numpy as np
    import sys
    sys.setrecursionlimit(10**7)
    from itertools import accumulate, combinations, permutations, product # https://docs.python.org/ja/3/library/itertools.html
    # accumulate() returns iterator! to get list: list(accumulate())
    from math import factorial, ceil, floor, sqrt
    def factorize(n):
        """return the factors of the Arg and count of each factor
        
        Args:
            n (long): number to be resolved into factors
        
        Returns:
            list of tuples: factorize(220) returns [(2, 2), (5, 1), (11, 1)]
        """
        fct = []  # prime factor
        b, e = 2, 0  # base, exponent
        while b * b <= n:
            while n % b == 0:
                n = n // b
                e = e + 1
            if e > 0:
                fct.append((b, e))
            b, e = b + 1, 0
        if n > 1:
            fct.append((n, 1))
        return fct
    def combinations_count(n, r):
        """Return the number of selecting r pieces of items from n kinds of items.
        
        Args:
            n (long): number
            r (long): number
        
        Raises:
            Exception: not defined when n or r is negative
        
        Returns:
            long: number
        """
        # TODO: How should I do when n - r is negative?
        if n < 0 or r < 0:
            raise Exception('combinations_count(n, r) not defined when n or r is negative')
        if n - r < r: r = n - r
        if r < 0: return 0
        if r == 0: return 1
        if r == 1: return n
        numerator = [n - r + k + 1 for k in range(r)]
        denominator = [k + 1 for k in range(r)]
        for p in range(2,r+1):
            pivot = denominator[p - 1]
            if pivot > 1:
                offset = (n - r) % p
                for k in range(p-1,r,p):
                    numerator[k - offset] /= pivot
                    denominator[k] /= pivot
        result = 1
        for k in range(r):
            if numerator[k] > 1:
                result *= int(numerator[k])
        return result
    def combinations_with_replacement_count(n, r):
        """Return the number of selecting r pieces of items from n kinds of items allowing individual elements to be repeated more than once.
        
        Args:
            n (long): number
            r (long): number
        
        Raises:
            Exception: not defined when n or r is negative
        
        Returns:
            long: number
        """
        if n < 0 or r < 0:
            raise Exception('combinations_with_replacement_count(n, r) not defined when n or r is negative')
        elif n == 0:
            return 1
        else:
            return combinations_count(n + r - 1, r)
    from bisect import bisect_left, bisect_right
    from collections import deque, Counter, defaultdict # https://docs.python.org/ja/3/library/collections.html#collections.deque
    from heapq import heapify, heappop, heappush, heappushpop, heapreplace,nlargest,nsmallest # https://docs.python.org/ja/3/library/heapq.html
    from copy import deepcopy, copy # https://docs.python.org/ja/3/library/copy.html
    import operator
    from operator import itemgetter #sort 
    # ex1: List.sort(key=itemgetter(1))
    # ex2: sorted(tuples, key=itemgetter(1,2))
    from functools import reduce, lru_cache
    def chmin(x, y):
        """change minimum
        if x > y, x = y and return (x, True).
        convenient when solving problems of dp[i]
        
        Args:
            x (long): current minimum value
            y (long): potential minimum value
        
        Returns:
            (x, bool): (x, True) when updated, else (x, False)
        """
        if x > y:
            x = y
            return (x, True)
        else:
            return (x, False)
    def chmax(x, y):
        """change maximum
        if x < y, x = y and return (x, True).
        convenient when solving problems of dp[i]
        
        Args:
            x (long): current maximum value
            y (long): potential maximum value
        
        Returns:
            (x, bool): (x, True) when updated, else (x, False)
        """
        if x < y:
            x = y
            return (x, True)
        else:
            return (x, False)

    from fractions import gcd # Deprecated since version 3.5: Use math.gcd() instead.
    def gcds(numbers):
        return reduce(gcd, numbers)
    def lcm(x, y):
        return (x * y) // gcd(x, y)
    def lcms(numbers):
        return reduce(lcm, numbers, 1)

    # first create factorial_list
    # fac_list = mod_factorial_list(n)
    INF = 10 ** 18
    MOD = 10 ** 9 + 7
    modpow = lambda a, n, p = MOD: pow(a, n, p) # Recursive function in python is slow!
    def modinv(a, p = MOD):
        # evaluate reciprocal using Fermat's little theorem:
        # a**(p-1) is identical to 1 (mod p) when a and p is coprime
        return modpow(a, p-2, p)
    def modinv_list(n, p = MOD):
        if n <= 1:
            return [0,1][:n+1]
        else:
            inv_t = [0,1]
            for i in range(2, n+1):
                inv_t += [inv_t[p % i] * (p - int(p / i)) % p]
            return inv_t
    def modfactorial_list(n, p = MOD):
        if n == 0:
            return [1]
        else:
            l = [0] * (n+1)
            tmp = 1
            for i in range(1, n+1):
                tmp = tmp * i % p
                l[i] = tmp
            return l
    def modcomb(n, k, fac_list = [], p = MOD):
        # fac_list = modfactorial_list(100)
        # print(modcomb(100, 5, modfactorial_list(100)))
        from math import factorial
        if n < 0 or k < 0 or n < k: return 0
        if n == 0 or k == 0: return 1
        if len(fac_list) <= n:
            a = factorial(n) % p
            b = factorial(k) % p
            c = factorial(n-k) % p
        else:
            a = fac_list[n]
            b = fac_list[k]
            c = fac_list[n-k]
        return (a * modpow(b, p-2, p) * modpow(c, p-2, p)) % p
    def modadd(a, b, p = MOD):
        return (a + b) % MOD
    def modsub(a, b, p = MOD):
        return (a - b) % p
    def modmul(a, b, p = MOD):
        return ((a % p) * (b % p)) % p
    def moddiv(a, b, p = MOD):
        return modmul(a, modpow(b, p-2, p))

    class UnionFindTree:
        """union find tree class
        TODO: fix this description...
        how to use (example):
        >>  uf = UnionFindTree(N) 
        >>  if uf.find_root(a) == uf.find_root(b):
        >>      do something
        >>  else:
        >>      do something
        >>      uf.unite(a, b)
        """
        def __init__(self, N):
            self.root = [-1] * (N+1)
            self.rank   = [0]  * (N+1)
            self.connected_num = [1] * (N+1)
    
        def find_root(self,x):
            root = self.root
            while root[x] != -1:
                x = root[x]
            return x
    
        def unite(self,x,y):
            root = self.root
            rank = self.rank
            connected_num = self.connected_num
            find_root = self.find_root
    
            rx = find_root(x)
            ry = find_root(y)
            if rx != ry:
                if rank[rx] < rank[ry]:
                    root[rx] = ry
                    rx,ry = ry,rx
                else:
                    if rank[rx] == rank[ry]:
                        rank[rx] += 1
                    root[ry] = rx       
                connected_num[rx] += connected_num[ry]

    """ initialize variables and set inputs
    # initialize variables
        # to initialize list, use [0] * n
        # to initialize two dimentional array, use [[0] * N for _ in range(N)]
    # set inputs
        # put inputs between specific values (almost as quickly)
        # ex) S = [-INF] + [int(r()) for _ in range(A)] + [INF]
        # open(0).read() is sometimes useful:
        # ex) n, m, *x = map(int, open(0).read().split())
        #     min(x[::2]) - max(x[1::2])
        # ex2) *x, = map(int, open(0).read().split())
        #     don't forget to add comma after *x if only one variable is used
    # preprocessing
        # transpose = [x for x in zip(*data)]
        # ex) [[1, 2, 3], [4, 5, 6], [7, 8, 9]] => [(1, 4, 7), (2, 5, 8), (3, 6, 9)]
        # flat = [flatten for inner in data for flatten in inner]
        # ex) [[1, 2, 3], [4, 5, 6], [7, 8, 9]] => [1, 2, 3, 4, 5, 6, 7, 8, 9]
    # calculate and output
        # output pattern
        # ex1) print(*l) => when l = [2, 5, 6], printed 2 5 6
    """

    # functions used
    r = lambda: sys.stdin.readline().strip()
    r_int = lambda: int(r())
    R = lambda: list(map(int, r().split()))
    Rmap = lambda: map(int, r().split())
    Rfloat = lambda: list(map(float, r().split()))
    Rtuple = lambda: tuple(map(int, r().split()))

    """ how to treat input
    # single int: int(r())
    # single string: r()
    # single float: float(r())
    # line int: R()
    # line string: r().split()
    # line (str, int, int): [j if i == 0 else int(j) for i, j in enumerate(r().split())]
    # lines int: [R() for _ in range(n)]
    """

    if __name__ != '__main__':
        # for test
        sys.stdin = open('sample.txt')
        
    # ----------------------------------
    # main

    N, K = R()
    TD = [Rtuple() for _ in range(N)]
    TD.sort(key=itemgetter(1), reverse=True)
    kind_sushi = set()
    num_kind = 0
    num_eaten = 0
    skip_candidate = []
    ans = 0
    tmp = 0
    for t, d in TD:
        if num_eaten == K:
            if t in kind_sushi or not skip_candidate:
                continue
            num_kind += 1
            tmp = tmp + d - skip_candidate.pop() + num_kind**2 - (num_kind-1)**2
            ans = max(ans, tmp)
        else:
            if t in kind_sushi:
                skip_candidate.append(d)
                num_eaten += 1
                tmp += d
            else:
                kind_sushi.add(t)
                num_kind += 1
                num_eaten += 1
                tmp += d
            if num_eaten == K:
                tmp = tmp + num_kind**2
                ans = tmp
    print(ans)
    
    # end of main
    # ----------------------------------

    """memo: how to use defaultdict of list
    # initialize
    Dic = defaultdict(list)
    # append / extend
    Dic[x].append(y)
    # three methods for loop: keys(), values(), items()
    for k, v in Dic.items():
    """
    
    """memo: how to solve binary problems
    # to make binary digits text
    >>> a = 5
    >>> bin_str_a = format(a, '#06b')
    >>> print(bin_str_a)
    0b0101 # first 2 strings (='0b') indicates it is binary
    """
    
    """memo: how to solve the problem
    creating simple test/answer
    greed
    simple dp
    graph
    """

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

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

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

# Step 1: Choose greedy
X.sort(key=lambda x: -x[1])
pq = []
chosen = defaultdict(int)
cnt = 0
score = 0
for t, d in X[:K]:
    heappush(pq, (d, t))
    cnt += 1 if chosen[t] == 0 else 0
    chosen[t] += 1
    score += d
ans = score + cnt ** 2

# Step 2: Replace chosen items
for t, d in X[K:]:
    if not pq:
        break
        
    cnt += 1 if chosen[t] == 0 else 0
    chosen[t] += 1
    score += d
    
    d_p, t_p = heappop(pq)
    chosen[t_p] -= 1
    cnt -= 1 if chosen[t_p] == 0 else 0
    score -= d_p
    ans = max(ans, score + cnt ** 2)
    
print(ans)

Submission status: Wrong Answer

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

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

# Step 1: Choose greedy
X.sort(key=lambda x: -x[1])
pq = []
chosen = defaultdict(int)
cnt = 0
score = 0
for t, d in X[:K]:
    heappush(pq, (d, t))
    cnt += 1 if chosen[t] == 0 else 0
    chosen[t] += 1
    score += d
ans = score + cnt ** 2

# Step 2: Replace chosen items
for t, d in X[K:]:
    if not pq:
        break
        
    chosen[t] += 1
    cnt += 1 if chosen[t] == 0 else 0
    score += d
    
    d_p, t_p = heappop(pq)
    chosen[t_p] -= 1
    cnt -= 1 if chosen[t_p] == 0 else 0
    score -= d_p
    ans = max(ans, score + cnt ** 2)
    
print(ans)

Submission status: Wrong Answer

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

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

# Step 1: Choose greedy
X.sort(key=lambda x: -x[1])
pq = []
chosen = defaultdict(int)
cnt = 0
score = 0
for t, d in X[:K]:
    heappush(pq, (d, t))
    cnt += 1 if chosen[t] == 0 else 0
    chosen[t] += 1
    score += d
ans = score + cnt ** 2

# Step 2: Replace chosen items
for t, d in X[K:]:
    if not pq:
        break
        
    if chosen[t] > 0:
        continue
        
    chosen[t] += 1
    cnt += 1
    d_p, t_p = heappop(pq)
    chosen[t_p] -= 1
    cnt -= 1 if chosen[t_p] == 0 else 0
    score += d - d_p
    ans = max(ans, score + cnt ** 2)
    
print(ans)

Submission status: Wrong Answer

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

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

# Step 1: Choose greedy
X.sort(key=lambda x: -x[1])
pq = []
chosen = defaultdict(int)
cnt = 0
score = 0
for t, d in X[:K]:
    heappush(pq, (d, t))
    cnt += 1 if chosen[t] == 0 else 0
    chosen[t] += 1
    score += d
ans = score + cnt ** 2

# Step 2: Replace chosen items
for t, d in X[K:]:
    if not pq:
        break
        
    if chosen[t] > 0:
        continue
        
    chosen[t] += 1
    cnt += 1
    d_p, t_p = heappop(pq)
    chosen[t_p] -= 1
    score += d - d_p
    ans = max(ans, score + cnt ** 2)
    
print(ans)

Submission status: Wrong Answer

Solution code:
from heapq import heappop, heappush
N, K = map(int, input().split())
X = [list(map(int, input().split())) for _ in range(N)]

# Step 1: Choose greedy
X.sort(key=lambda x: -x[1])
pq = []
chosen = [0] * (N + 1)
score = 0
for t, d in X[:K]:
    heappush(pq, (d, t))
    chosen[t] += 1
    score += d
ans = score + sum(x > 0 for x in chosen) ** 2

# Step 2: Replace chosen items
for t, d in X[K:]:
    if not pq:
        break
        
    if chosen[t] > 0:
        continue
        
    chosen[t] += 1
    d_p, t_p = heappop(pq)
    chosen[t_p] -= 1
    score += d - d_p
    ans = max(ans, score + sum(x > 0 for x in chosen) ** 2)
    
print(ans)

Submission status: Wrong Answer

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

tds.sort(key=lambda x:x[1])

arr = {}
for idx, td in enumerate(tds):
    t, d = td
    arr[t] = [d, n - idx - 1]

tds = [y for x, y in tds[::-1]]

base = 0
rest = sum(tds[:k])
res = 0
kind = 0
new_arr = sorted(arr.values(), reverse=True)

for x, y in new_arr:
    kind += 1
    base += x
    if y < k:
        rest -= x
        tds[y] = 0
    else:
        while k >= 1 and tds[k - 1] == 0:
            k -= 1
        rest -= tds[k - 1]
        k -= 1
        if k <= 0:
            break
    if k >= 0:
        res = max(res, kind ** 2 + base + rest)

print(res)