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 heapq

from typing import List


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

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        result = max(result, op + types ** 2)

    return result


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import heapq

from typing import List


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

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(key=lambda x: x[1], reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import heapq

from typing import List


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

    print(vs(v, k))


def vs(v: List[List[int]], k: int) -> int:
    v.sort(reverse=True)

    # Dup type
    h0 = []
    # New type
    h1 = []
    # types
    s = set()
    types = 0
    op = 0

    for i in range(k):
        t, d = v[i]
        op += d

        if t in s:
            heapq.heappush(h0, d)
            continue

        types += 1
        s.add(t)

    for i in range(k, len(v)):
        t, d = v[i]

        if not t in s:
            heapq.heappush(h1, -d)

    result = op + types ** 2

    while len(h0) > 0 and len(h1) > 0:
        d0 = heapq.heappop(h0)
        d1 = -heapq.heappop(h1)
        op = op - d0 + d1
        types += 1
        tmp = op + types ** 2

        if tmp > result:
            result = tmp

    return result


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import heapq
from itertools import accumulate
n,k=map(int,input().split())

td=[[] for _ in range(n+1)]
kind=[]
base=[]
used=[0]*(n+1)

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

    heapq.heappush(td[t],-d)

for i in range(1,n+1):
    if len(td[i])==0:
        continue

    kind.append(-heapq.heappop(td[i]))
    
    while len(td[i])>0:
        base.append(-heapq.heappop(td[i]))

kind.sort(reverse=True)
base.sort(reverse=True)
kind=list(accumulate(kind))
base=list(accumulate(base))

ans=0

for kindcnt in range(max(1,k-len(base)),len(kind)+1):
    basecnt=k-kindcnt
    tmp=kind[kindcnt-1]+kindcnt**2
    if basecnt>0:
        tmp+=base[basecnt-1]
    ans=max(ans,tmp)

print(ans)

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(10**9)

def mi(): return map(int,input().split())
def ii(): return int(input())
def isp(): return input().split()
def deb(text): print("-------\n{}\n-------".format(text))

class Counter:
    def __init__(self):
        self.dict = {}

    def add(self,x):
        if x in self.dict: self.dict[x] += 1
        else: self.dict[x] = 1
    
    def decrement(self,x):
        self.dict[x] -= 1


    def get_dict(self):
        return self.dict



INF=10**20
def main():
    N,K=mi()
    X = []
    for i in range(N):
        t,d = mi()
        X.append((t,d))
    
    X.sort(key=lambda  x: x[1],reverse=True)

    base = X[:K][::-1]
    rest = X[K:]

    base_sum = 0
    seen = set()
    counter = Counter()
    for t,d in base:
        seen.add(t)
        base_sum += d
        counter.add(t)

    base_kind = len(seen)
    base_sum += base_kind ** 2

    ans = base_sum
    i = 0
    for j in range(len(rest)):
        if not i < K: continue
        rt,rd = rest[j] # 足していく
        bt,bd = base[i] # 消していく
        
        if rt in seen: continue
        dic = counter.get_dict()
        # rt_count = 0
        # if rd in dic:
        #     rt_count = dic[rd]
            
        bt_count = dic[bt]

        x = len(seen)
        if bt_count == 1:
            seen.remove(bt)

        counter.decrement(bt)
        counter.add(rt)
        seen.add(rt)
        new_x = len(seen)
        if new_x <= x:
            seen.add(bt)
            counter.add(bt)
            counter.decrement(rt)
            if counter.get_dict()[rt] == 0: seen.remove(rt)

            continue

        base_sum = base_sum - bd + rd - x**2 + new_x ** 2
        ans = max(ans,base_sum)

        i += 1
    
    print(ans)
        


                




if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(10**9)

def mi(): return map(int,input().split())
def ii(): return int(input())
def isp(): return input().split()
def deb(text): print("-------\n{}\n-------".format(text))

class Counter:
    def __init__(self):
        self.dict = {}

    def add(self,x):
        if x in self.dict: self.dict[x] += 1
        else: self.dict[x] = 1
    
    def decrement(self,x):
        self.dict[x] -= 1


    def get_dict(self):
        return self.dict



INF=10**20
def main():
    N,K=mi()
    X = []
    for i in range(N):
        t,d = mi()
        X.append((t,d))
    
    X.sort(key=lambda  x: x[1],reverse=True)

    base = X[:K][::-1]
    rest = X[K:]

    base_sum = 0
    seen = set()
    counter = Counter()
    for t,d in base:
        seen.add(t)
        base_sum += d
        counter.add(t)

    base_kind = len(seen)
    base_sum += base_kind ** 2

    ans = base_sum
    j = 0
    for i in range(K):
        if not j < len(rest): break
        rt,rd = rest[j] # 足していく
        bt,bd = base[i] # 消していく
        
        if rt in seen: j += 1; continue
        dic = counter.get_dict()
        # rt_count = 0
        # if rd in dic:
        #     rt_count = dic[rd]
            
        bt_count = dic[bt]

        x = len(seen)
        if bt_count == 1:
            seen.remove(bt)

        counter.decrement(bt)
        counter.add(rt)
        seen.add(rt)
        new_x = len(seen)

        base_sum = base_sum - bd + rd - x**2 + new_x ** 2
        ans = max(ans,base_sum)

        j += 1
    
    print(ans)
        


                




if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(10**9)

def mi(): return map(int,input().split())
def ii(): return int(input())
def isp(): return input().split()
def deb(text): print("-------\n{}\n-------".format(text))

class Counter:
    def __init__(self):
        self.dict = {}

    def add(self,x):
        if x in self.dict: self.dict[x] += 1
        else: self.dict[x] = 1
    
    def decrement(self,x):
        self.dict[x] -= 1


    def get_dict(self):
        return self.dict



INF=10**20
def main():
    N,K=mi()
    X = []
    for i in range(N):
        t,d = mi()
        X.append((t,d))
    
    X.sort(key=lambda  x: x[1],reverse=True)

    base = X[:K][::-1]
    rest = X[K:]

    base_sum = 0
    seen = set()
    counter = Counter()
    for t,d in base:
        seen.add(t)
        base_sum += d
        counter.add(t)

    base_kind = len(seen)
    base_sum += base_kind ** 2

    ans = base_sum
    for i in range(min(len(base),len(rest))):
        rt,rd = rest[i] # 足していく
        bt,bd = base[i] # 消していく

        dic = counter.get_dict()
        rt_count = 0
        if rd in dic:
            rt_count = dic[rd]
            
        bt_count = dic[bt]

        x = len(seen)
        if bt_count == 1:
            seen.remove(bt)

        counter.decrement(bt)
        counter.add(rt)
        seen.add(rt)
        new_x = len(seen)

        base_sum = base_sum - bd + rd - x**2 + new_x ** 2
        ans = max(ans,base_sum)
    
    print(ans)
        


                




if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(10**9)

def mi(): return map(int,input().split())
def ii(): return int(input())
def isp(): return input().split()
def deb(text): print("-------\n{}\n-------".format(text))

class Counter:
    def __init__(self):
        self.dict = {}

    def add(self,x):
        if x in self.dict: self.dict[x] += 1
        else: self.dict[x] = 1
    
    def decrement(self,x):
        self.dict[x] -= 1


    def get_dict(self):
        return self.dict



INF=10**20
def main():
    N,K=mi()
    X = []
    for i in range(N):
        t,d = mi()
        X.append((t,d))
    
    X.sort(key=lambda  x: x[1],reverse=True)

    base = X[:K][::-1]
    rest = X[K:]

    base_sum = 0
    seen = set()
    counter = Counter()
    for t,d in base:
        seen.add(t)
        base_sum += d
        counter.add(t)

    base_kind = len(seen)
    base_sum += base_kind ** 2

    ans = base_sum
    for i in range(min(len(base),len(rest))):
        rt,rd = rest[i] # 足していく
        bt,bd = base[i] # 消していく

        dic = counter.get_dict()
        rt_count = 0
        if rd in dic:
            rt_count = dic[rd]
            
        bt_count = dic[bt]

        x = len(seen)
        if bt_count == 1:
            counter.decrement(bt)
            seen.remove(bt)
        
        counter.add(rt)
        seen.add(rt)
        new_x = len(seen)

        base_sum = base_sum - bd + rd - x**2 + new_x ** 2
        ans = max(ans,base_sum)
    
    print(ans)
        


                




if __name__ == "__main__":
    main()

Submission status: Wrong Answer

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

def main():
  N, K=map(int,input().split())
  sushies=[list(map(int,input().split())) for i in range(N)]
  sushies.sort(reverse=True, key=lambda x:x[1])
  dsum=0
  tset=set()
  kubi=[]
  for i in range(K):
    dsum+=sushies[i][1]
    t=sushies[i][0]
    if t in tset:
      kubi.append(sushies[i])
    tset.add(t)
  x=len(tset)
  koho=[]
  for i in range(K,N):
    t=sushies[i][0]
    if not(t in tset):
      koho.append(sushies[i])
  ind=0
  ans=dsum+x*x
  while(ind < len(kubi) and ind < len(koho)):
    dsum=dsum-kubi[-1-ind][1]+koho[ind][1]
    x+=1
    ans=max(ans, dsum+x*x)
    ind+=1
  print(ans)
  
if __name__=="__main__":
  main()

Submission status: Wrong Answer

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

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

data.sort(key=lambda x:-x[1])
KISO = 0
EAT = defaultdict(int)
for i in range(K):
  t,v = data[i]
  EAT[t] += 1
  KISO += v
SHURUI = len(EAT)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if EAT[a] == 1:
    l -= 1
    continue
  if EAT[c] != 0:
    r += 1
    continue
  EAT[a] -= 1
  EAT[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

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

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

data.sort(lambda x:-x[1])
KISO = 0
EAT = defaultdict(int)
for i in range(K):
  t,v = data[i]
  EAT[t] += 1
  KISO += v
SHURUI = len(EAT)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if EAT[a] == 1:
    l -= 1
    continue
  if EAT[c] != 0:
    r += 1
    continue
  EAT[a] -= 1
  EAT[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

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

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

data.sort(lambda x:-x[1])
KISO = 0
TOP = defaultdict(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if TOP[a] < 2:
    l -= 1
    continue
  if TOP[c] != 0:
    r += 1
    continue
  TOP[a] -= 1
  TOP[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict

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

data.sort(lambda x:-x[1])
KISO = 0
TOP = defaultdict(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if TOP[a] == 1:
    l -= 1
    continue
  if TOP[c] != 0:
    r += 1
    continue
  TOP[a] -= 1
  TOP[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

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

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache

#############
# Constants #
#############

MOD = 10**9+7
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)

#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if n == r: return 1
  if n < r or r < 0: return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1: arr.append([temp, 1])
  if arr==[]: arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b: a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n-1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n: return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

def base_10_to_n_without_0(X, n):
  X -= 1
  if X//n: return base_10_to_n_without_0(X//n, n)+[X%n]
  return [X%n]

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############

N,K = IL()
data = ILL(N)

data.sort(lambda x:-x[1])
KISO = 0
TOP = DD(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  a,b = data[l]
  c,d = data[r]
  if TOP[a] == 1:
    l -= 1
    continue
  if TOP[c] != 0:
    r += 1
    continue
  TOP[a] -= 1
  TOP[c] += 1
  KISO -= b
  KISO += d
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

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

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache

#############
# Constants #
#############

MOD = 10**9+7
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)

#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if n == r: return 1
  if n < r or r < 0: return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1: arr.append([temp, 1])
  if arr==[]: arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b: a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n-1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n: return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

def base_10_to_n_without_0(X, n):
  X -= 1
  if X//n: return base_10_to_n_without_0(X//n, n)+[X%n]
  return [X%n]

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############

N,K = IL()
data = ILL(N)

data.sort(lambda x:-x[1])
KISO = 0
TOP = DD(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N:
    break
  elif TOP[data[l][0]] == 1:
    l -= 1
  elif TOP[data[r][0]] != 0:
    r += 1
  else:
    TOP[data[l][0]] -= 1
    TOP[data[r][0]] += 1
    KISO -= data[l][1]
    KISO += data[r][1]
    SHURUI += 1
    MAX = max(MAX,KISO+SHURUI**2)
print(MAX)

Submission status: Wrong Answer

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

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache

#############
# Constants #
#############

MOD = 10**9+7
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)

#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if n == r: return 1
  if n < r or r < 0: return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1: arr.append([temp, 1])
  if arr==[]: arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b: a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n-1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n: return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

def base_10_to_n_without_0(X, n):
  X -= 1
  if X//n: return base_10_to_n_without_0(X//n, n)+[X%n]
  return [X%n]

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############

N,K = IL()
data = ILL(N)

data.sort(lambda x:-x[1])
KISO = 0
TOP = DD(int)
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N:
    print(MAX)
    exit()
  if TOP[data[l][0]] == 1:
    l -= 1
    continue
  if TOP[data[r][0]] != 0:
    r += 1
    continue
  TOP[data[l][0]] -= 1
  TOP[data[r][0]] += 1
  KISO -= data[l][1]
  KISO += data[r][1]
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
#print(MAX)

Submission status: Wrong Answer

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

#############
# Libraries #
#############

import sys
input = sys.stdin.readline

import math
#from math import gcd
import bisect
import heapq
from collections import defaultdict
from collections import deque
from collections import Counter
from functools import lru_cache

#############
# Constants #
#############

MOD = 10**9+7
INF = float('inf')
AZ = "abcdefghijklmnopqrstuvwxyz"

#############
# Functions #
#############

######INPUT######
def I(): return int(input().strip())
def S(): return input().strip()
def IL(): return list(map(int,input().split()))
def SL(): return list(map(str,input().split()))
def ILs(n): return list(int(input()) for _ in range(n))
def SLs(n): return list(input().strip() for _ in range(n))
def ILL(n): return [list(map(int, input().split())) for _ in range(n)]
def SLL(n): return [list(map(str, input().split())) for _ in range(n)]

######OUTPUT######
def P(arg): print(arg); return
def Y(): print("Yes"); return
def N(): print("No"); return
def E(): exit()
def PE(arg): print(arg); exit()
def YE(): print("Yes"); exit()
def NE(): print("No"); exit()

#####Shorten#####
def DD(arg): return defaultdict(arg)

#####Inverse#####
def inv(n): return pow(n, MOD-2, MOD)

######Combination######
kaijo_memo = []
def kaijo(n):
  if(len(kaijo_memo) > n): return kaijo_memo[n]
  if(len(kaijo_memo) == 0): kaijo_memo.append(1)
  while(len(kaijo_memo) <= n): kaijo_memo.append(kaijo_memo[-1] * len(kaijo_memo) % MOD)
  return kaijo_memo[n]

gyaku_kaijo_memo = []
def gyaku_kaijo(n):
  if(len(gyaku_kaijo_memo) > n): return gyaku_kaijo_memo[n]
  if(len(gyaku_kaijo_memo) == 0): gyaku_kaijo_memo.append(1)
  while(len(gyaku_kaijo_memo) <= n): gyaku_kaijo_memo.append(gyaku_kaijo_memo[-1] * pow(len(gyaku_kaijo_memo),MOD-2,MOD) % MOD)
  return gyaku_kaijo_memo[n]

def nCr(n,r):
  if n == r: return 1
  if n < r or r < 0: return 0
  ret = 1
  ret = ret * kaijo(n) % MOD
  ret = ret * gyaku_kaijo(r) % MOD
  ret = ret * gyaku_kaijo(n-r) % MOD
  return ret

######Factorization######
def factorization(n):
  arr = []
  temp = n
  for i in range(2, int(-(-n**0.5//1))+1):
    if temp%i==0:
      cnt=0
      while temp%i==0: 
        cnt+=1 
        temp //= i
      arr.append([i, cnt])
  if temp!=1: arr.append([temp, 1])
  if arr==[]: arr.append([n, 1])
  return arr

#####MakeDivisors######
def make_divisors(n):
  divisors = []
  for i in range(1, int(n**0.5)+1):
    if n % i == 0:
      divisors.append(i)
      if i != n // i: 
        divisors.append(n//i)
  return divisors

#####MakePrimes######
def make_primes(N):
  max = int(math.sqrt(N))
  seachList = [i for i in range(2,N+1)]
  primeNum = []
  while seachList[0] <= max:
    primeNum.append(seachList[0])
    tmp = seachList[0]
    seachList = [i for i in seachList if i % tmp != 0]
  primeNum.extend(seachList)
  return primeNum

#####GCD#####
def gcd(a, b):
    while b: a, b = b, a % b
    return a

#####LCM#####
def lcm(a, b):
    return a * b // gcd (a, b)

#####BitCount#####
def count_bit(n):
  count = 0
  while n:
    n &= n-1
    count += 1
  return count

#####ChangeBase#####
def base_10_to_n(X, n):
  if X//n: return base_10_to_n(X//n, n)+[X%n]
  return [X%n]

def base_n_to_10(X, n):
  return sum(int(str(X)[-i-1])*n**i for i in range(len(str(X))))

def base_10_to_n_without_0(X, n):
  X -= 1
  if X//n: return base_10_to_n_without_0(X//n, n)+[X%n]
  return [X%n]

#####IntLog#####
def int_log(n, a):
  count = 0
  while n>=a:
    n //= a
    count += 1
  return count

#############
# Main Code #
#############

N,K = IL()
data = ILL(N)

data.sort(lambda x:-x[1])
KISO = 0
TOP = [0 for i in range(N+1)]
for i in range(K):
  t,v = data[i]
  TOP[t] += 1
  KISO += v
SHURUI = 0
for i in range(N+1):
  if TOP[i]!=0:
    SHURUI+=1
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  if TOP[data[l][0]] == 1:
    l -= 1
    continue
  if TOP[data[r][0]] != 0:
    r += 1
    continue
  TOP[data[l][0]] -= 1
  TOP[data[r][0]] += 1
  KISO -= data[l][1]
  KISO += data[r][1]
  SHURUI += 1
  MAX = max(MAX,KISO+SHURUI**2)
print(MAX)