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:
# -*- 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
  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.append(KISO+SHURUI**2)
print(max(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
  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])
dic = DD(list)
for t,v in data:
  dic[t].append(v)

KISO = 0
TOP = DD(list)
for i in range(K):
  t,v = data[i]
  TOP[t].append(v)
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  if len(TOP[data[l][0]]) == 1:
    l -= 1
    continue
  if len(TOP[data[r][0]]) != 0:
    r += 1
    continue
  TOP[data[l][0]].pop()
  TOP[data[r][0]].append(data[r][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])
dic = DD(list)
for t,v in data:
  dic[t].append(v)

KISO = 0
TOP = DD(list)
for i in range(K):
  t,v = data[i]
  TOP[t].append(v)
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  if len(TOP[data[l][0]]) == 1:
    l -= 1
    continue
  if TOP[data[r][0]]:
    r += 1
    continue
  TOP[data[l][0]].pop()
  TOP[data[r][0]].append(data[r][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])
dic = DD(list)
for t,v in data:
  dic[t].append(v)
for t in dic:
  dic[t].sort(reverse=True)

KISO = 0
TOP = DD(list)
for i in range(K):
  t,v = data[i]
  TOP[t].append(v)
  KISO += v
SHURUI = len(TOP)
l = K-1
r = K
MAX = KISO+SHURUI**2
while True:
  if l<0 or r>=N: break
  if len(TOP[data[l][0]]) == 1:
    l -= 1
    continue
  if TOP[data[r][0]]:
    r += 1
    continue
  TOP[data[l][0]].pop()
  TOP[data[r][0]].append(data[r][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 -*-

import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [0], [0]
    st = set()

    for t, d in S:
        if t in st:
            #累積報酬和のリスト
            y_0.append(y_0[-1]+d)
        else:
            st.add(t)
            y_1.append(y_1[-1]+d)
    

    max_sum = -1

    for i in range(1,k+1):
        print(i)
        if k-i >= len(y_0):
            continue
        if i >= len(y_1):
            break
        max_sum = max(max_sum, y_0[k-i]+y_1[i]+i**2)
    
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    S = [tuple(map(int, input().split())) for _ in range(n)]
    S.sort(key=lambda x: -x[1])
    y_1, y_0 = [], []
    st = set()
    for t, d in S:
        if t in st:
            y_0.append(d)
        else:
            st.add(t)
            y_1.append(d)

    max_sum = -1
    for number_of_y_1 in range(len(y_1)+1):
        current_sum = 0
        number_of_y_0 = k - number_of_y_1
        current_sum += sum(y_1[:number_of_y_1])
        current_sum += sum(y_0[:number_of_y_0])
        current_sum += number_of_y_1 ** 2
        if max_sum == -1:
            max_sum = current_sum
        else:
            if max_sum < current_sum:
                max_sum = current_sum
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import sys

def main():
    input = sys.stdin.readline
    n, k = map(int, input().split())
    susi_dict = {}

    for i in range(n):
        t, d = map(int, input().split())
        if t in susi_dict:
            susi_dict[t].append(d)
        else:
            susi_dict[t] = []
            susi_dict[t].append(d)
    
    y_1 = []
    y_0 = []
    for susi_list in susi_dict.values():
        list.sort(susi_list, reverse=True)
        y_1.append(susi_list[0])
        y_0.extend(susi_list[1:])
    list.sort(y_1, reverse=True)
    list.sort(y_0, reverse=True)

    max_sum = -1
    for number_of_y_1 in range(len(y_1)):
        current_sum = 0
        number_of_y_0 = k - number_of_y_1
        current_sum += sum(y_1[:number_of_y_1])
        current_sum += sum(y_0[:number_of_y_0])
        current_sum += number_of_y_1 ** 2
        if max_sum == -1:
            max_sum = current_sum
        else:
            if max_sum < current_sum:
                max_sum = current_sum
    print(max_sum)
    return


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import sys
def input():return sys.stdin.readline()[:-1]
def N(): return int(input())
def NM():return map(int,input().split())
def L():return list(NM())
def LN(n):return [N() for i in range(n)]
def LL(n):return [L() for i in range(n)]
n,k=NM()
l=LL(n)
l.sort(key=lambda x:-x[1])
from heapq import heappush,heappop,heapify
s=set()
q=[]
oisisa=0
syu=0
for t,d in l[:k]:
    if t in s:
        heappush(q,d)
    else:
        s.add(t)
        syu+=1
    oisisa+=d
ans=oisisa+syu**2
for t,d in l[k:]:
    if q and not t in s:
        syu+=1
        oisisa+=d-heappop(q)
    ans=max(ans,oisisa+syu**2)
print(ans)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mi = []
	sushi_ki = []
	syu = {}
	for i in range(N):
		t,d = inputMap()
		if t in syu:
			sushi_ki.append(d)
		else:
			syu[t] = 1
			sushi_mi.append(d)

	sushi_mi.sort()
	sushi_mi = sushi_mi[::-1]
	sushi_ki.sort()
	sushi_ki = sushi_ki[::-1]

	ans = 0
	for i in range(1, min(len(sushi_mi), K+1)+1):
		x = 0
		tmp = 0
		if K - i > len(sushi_ki):
			continue

		for j in range(0, i):
			tmp += sushi_mi[j]
			x += 1

		for j in range(0, K - i):
			tmp += sushi_ki[j]

		oishi = tmp + x * x
		if ans < oishi:
			ans = oishi
	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
import sys
import math
from bisect import bisect_left
from bisect import bisect_right
from collections import defaultdict
from heapq import heappop, heappush
import itertools
from decimal import *

input = sys.stdin.readline

def inputInt(): return int(input())
def inputMap(): return map(int, input().split())
def inputList(): return list(map(int, input().split()))
def inputStr(): return input()[:-1]

inf = float('inf')
mod = 1000000007

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-

def main():
	N,K = inputMap()
	sushi_mi = []
	sushi_ki = []
	syu = {}
	for i in range(N):
		t,d = inputMap()
		if t in syu:
			sushi_ki.append(d)
		else:
			syu[t] = 1
			sushi_mi.append(d)

	sushi_mi.sort()
	sushi_mi = sushi_mi[::-1]
	sushi_ki.sort()
	sushi_ki = sushi_ki[::-1]

	ans = 0
	for i in range(1,K+1):
		if i > len(sushi_mi):
			continue
		if K-i > len(sushi_ki):
			continue
			
		tmp = 0
		for j in range(0,i):
			tmp += sushi_mi[j]
		for j in range(0,K-i):
			tmp += sushi_ki[j]

		tmp += i*i
		if tmp > ans:
			ans = tmp

	print(ans)

#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
#-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-
if __name__ == "__main__":
	main()

Submission status: Wrong Answer

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

dset = [[] for i in range(n)]
for i in range(n):
    t, d = map(int, input().split())
    dset[t-1].append(d)

for dd in dset:
    if dd: dd.sort(reverse=True)
    else: dd.append(-float('inf'))

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

cur = 0
hq = []
# k種類の先頭を除いて格納
for i in range(k):
    cur+=dset[i][0]
    for d in dset[i][1:]: heapq.heappush(hq, -d)

# 残りの種類のすべてを格納
for dd in dset[k:]:
    for d in dd:
        heapq.heappush(hq, -d)

ans = cur + k*k
for i in range(k-1, -1, -1):
    v = dset[i][0] # 削除候補
    if hq:
        w = -heapq.heappop(hq)
    else:
        break
    if v < w:
        cur += w - v
    else:
        heapq.heappush(hq, -w)
    ans = max(ans, cur + i*i)
print(ans)

Submission status: Wrong Answer

Solution code:
n,k= map(int,input().split())
tlist=[]
dlist=[]
fulllist = []
for i in range(n):
    t,d= map(int,input().split())
    tlist.append(t)
    dlist.append(d)
    fulllist.append((d,t))
newflist = sorted(fulllist,reverse=True)
basesum = 0
baseadd = []
shurui = set()
basedict = {}
inbase = []
for (c,e) in sorted(newflist[:k],reverse=True):
    basesum+=c
    baseadd.append((c,e))
    if e in shurui:        
        inbase.append((c,e))
    else:
        shurui.add(e)
    basedict[e]=basedict.get(e,[])+[c]
basepoint = basesum + len(shurui)*len(shurui)
if n==k:
    print(basepoint)
else:
    outbase = []
    outset = set()
    for (c,e) in newflist[k:]:
        if (not e in shurui) and (not e in outset):
            outbase.append((c,e))
            outset.add(e)
    koukan = min(len(inbase),len(outbase))
    points=[basepoint]
    g=basepoint
    sh = len(shurui)
    for kout,kin in zip(sorted(outbase[:koukan],reverse=True),sorted(inbase[:koukan])):
        newpoint = points[-1]-kin[0]+kout[0] + 2*sh + 1
        sh+=1
        points.append(newpoint)
        g = max(g,newpoint)
    print(g)

Submission status: Wrong Answer

Solution code:
n,k= map(int,input().split())
tlist=[]
dlist=[]
fulllist = []
for i in range(n):
    t,d= map(int,input().split())
    tlist.append(t)
    dlist.append(d)
    fulllist.append((d,t))
newflist = sorted(fulllist,reverse=True)
basesum = 0
baseadd = []
shurui = set()
basedict = {}
inbase = []
for (c,e) in sorted(newflist[:k],reverse=True):
    basesum+=c
    baseadd.append((c,e))
    if e in shurui:        
        inbase.append((c,e))
    else:
        shurui.add(e)
    basedict[e]=basedict.get(e,[])+[c]
basepoint = basesum + len(shurui)*len(shurui)
if n==k:
    print(basepoint)
else:
    outbase = []
    outset = set()
    for (c,e) in newflist[k:]:
        if (not e in shurui) and (not e in outset):
            outbase.append((c,e))
            outset.add(e)
    koukan = min(len(inbase),len(outbase))
    points=[basepoint]
    g=basepoint
    sh = len(shurui)
    for kout,kin in zip(outbase[:koukan],sorted(inbase[:koukan])):
        newpoint = points[-1]-kin[0]+kout[0] + 2*sh + 1
        sh+=1
        points.append(newpoint)
        g = max(g,newpoint)
    print(g)