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: Runtime Error

Solution code:
import sys
from collections import deque
from heapq import heappop, heappush

sys.setrecursionlimit(10 ** 7)
input = sys.stdin.readline
f_inf = float('inf')
mod = 10 ** 9 + 7


def resolve():
    n, k = map(int, input().split())
    LD = [[] for _ in range(n)]
    for _ in range(n):
        t, d = map(int, input().split())
        LD[t - 1].append([d, 0] if len(LD[t - 1]) else [d, 1])

    ALL = []
    for ld in LD:
        ALL.extend(ld)
    ALL.sort(reverse=True)
    que = deque(ALL)

    x = y = 0
    zero, one = [], []
    for _ in range(k):
        d, ex = que.popleft()
        x += d
        y += ex
        heappush(zero, d) if ex == 0 else heappush(one, d)

    res = x + y ** 2
    while que:
        d, ex = que.popleft()
        if ex:
            if zero:
                x += d - heappop(zero)
                y += 1
                res = max(res, x + y ** 2)
            else:
                x += d - heappop(one)
                res = max(res, x + y ** 2)
    print(res)


if __name__ == '__main__':
    resolve()

Submission status: Runtime Error

Solution code:
from heapq import heapify, heappop
from sys import stdin
def input():
    return stdin.readline().strip()

def main():
    n, k = map(int, input().split())

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

    dt.sort(key=lambda x: x[0], reverse=True)

    # select k biggest points
    base = 0
    variety = [0] * n
    bonus = 0
    heap = []
    for d, t in dt[:k]:
        base += d
        if variety[t] == 0:
            bonus += 1
        variety[t] += 1
        heap.append((d, t))
    heapify(heap)

    # select exchange candidates
    new_dt = []
    for d, t in dt[k:]:
        if variety[t] == 0:
            variety[t] = 1
            new_dt.append((d, t))

    # exchange
    ans = base + bonus * bonus
    for new_d, new_t in new_dt:
        if len(heap) == 1:
            break
        d, t = heappop(heap)
        while len(heap) > 0 and variety[t] == 1:
            d, t = heappop(heap)
        variety[t] -= 1

        base -= d - new_d
        bonus += 1
        if ans < base + bonus * bonus:
            ans = base + bonus * bonus

    print(ans)

main()

Submission status: Runtime Error

Solution code:
from heapq import heapify, heappop
from sys import stdin
def input():
    return stdin.readline().strip()

def main():
    n, k = map(int, input().split())

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

    dt.sort(key=lambda x: x[0], reverse=True)

    # select k biggest points
    base = 0
    variety = [0] * n
    bonus = 0
    heap = []
    for d, t in dt[:k]:
        base += d
        if variety[t] == 0:
            bonus += 1
        variety[t] += 1
        heap.append((d, t))
    heapify(heap)

    # select exchange candidates
    new_dt = []
    for d, t in dt[k:]:
        if variety[t] == 0:
            variety[t] = 1
            new_dt.append((d, t))

    # exchange
    ans = base + bonus * bonus
    for new_d, new_t in new_dt:
        d, t = heappop(heap)
        while len(heap) > 0 and variety[t] == 1:
            d, t = heappop(heap)
        variety[t] -= 1

        base -= d - new_d
        bonus += 1
        if ans < base + bonus * bonus:
            ans = base + bonus * bonus

    print(ans)

main()

Submission status: Runtime Error

Solution code:
import heapq
from sys import stdin

N,K = map(int,input().split())

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = 0

if K >= len(sushi.keys()):
    max_variation = K
else:
    max_variation = len(sushi.keys())
    for _ in range(K-max_variation):
        cur += heapq.heappop(Q)

cur += sum(P)
ans = cur + max_variation**2

for i in range(max_variation-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(max_variation-i-1)**2)

print(ans)

Submission status: Runtime Error

Solution code:
import heapq
from sys import stdin

N,K = map(int,input().split())

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = sum(P)
ans = cur + K**2

for i in range(max(K,len(sushi.keys()))-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(K-i-1)**2)

print(ans)

Submission status: Runtime Error

Solution code:
import heapq
from sys import stdin

N,K = map(int,input().split())

sushi = {}

for i in range(N):
    t,d = map(int, stdin.readline().split())
    if t in sushi.keys():
        sushi[t].append(d)
    else:
        sushi[t] = [d]

P = []
Q = []

for x in sushi.keys():
    sushi[x].sort(reverse = True)
    P.append(sushi[x][0])
    
    for i in range(1,len(sushi[x])):
        Q.append(sushi[x][i])

heapq.heapify(P)

Q = list(map(lambda x:-x,Q))
heapq.heapify(Q)

while len(P) > K:
    m = heapq.heappop(P)
    heapq.heappush(Q,-m)

cur = sum(P)
ans = cur + K**2

for i in range(K-1):
    p = heapq.heappop(P)
    cur -= p
    q = -heapq.heappushpop(Q,-p)
    cur += q
    ans = max(ans,cur+(K-i-1)**2)

print(ans)

Submission status: Runtime Error

Solution code:
import heapq

n, k = map(int, input().split())
a = [0] * (n + 1)
typecnt = 0
h = []
heapq.heapify(h)
for _ in range(n):
    t, d = map(int, input().split())
    if a[t] == 0:
        typecnt += 1
        a[t] = 1
    heapq.heappush(h, [-d, t])
ans = 0
mintype = 0
m = []
heapq.heapify(m)
for _ in range(k):
    d, t = heapq.heappop(h)
    if a[t] == 1:
        mintype += 1
        a[t] = 0
    else:
        heapq.heappush(m, -d)
    ans -= d
cnt = ans
ans += mintype * mintype
for j in range(mintype, typecnt):
    bonus = (j + 1) * (j + 1)
    while h:
        d, t = heapq.heappop(h)
        if a[t] == 1:
            a[t] = 0
            l = heapq.heappop(m)
            cnt -= (d + l)
            ans = max(ans, cnt + bonus)
            break
print(ans)

Submission status: Runtime Error

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

for _ in range(N):
  TD.append(tuple(map(int, input().split())))

TD.sort(key = lambda x:x[1], reverse=True)
TD += [(TD[-1][-1], 0)]

ans1 = 0
T1 = [0]*N
L = []
cnt1 = 0
temp = TD[0][1]

ans2 = 0
T2 = [0]*N
D2 = []
cnt2 = 0

for t, d in TD:
  if cnt1 < K:
    if temp == d and T1[t-1] == 1:
      L.append((t, d))
    elif temp == d and T1[t-1] == 0:
      T1[t-1] = 1
      ans1 += d
      cnt1 += 1
    else:
      while L:
        if cnt1 == K:
          break
        a, b = L.pop()
        ans1 += b
        cnt1 += 1
      if cnt1 < K:
        T1[t-1] = 1
        ans1 += d
        cnt1 += 1
        temp = d
  if cnt2 < K:
    if T2[t-1]:
      D2.append(d)
    else:
      T2[t-1] = 1
      ans2 += d
      cnt2 += 1

ans1 += sum(T1)**2
ans2 += sum(D2[:K-cnt2]) + sum(T2)**2

print(max(ans1, ans2))

Submission status: Runtime Error

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)
        assert - x**2 + new_x ** 2 != 0

        j += 1
    
    print(ans)
        


                




if __name__ == "__main__":
    main()

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((t-1, d))

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

import heapq

Tcount = [0]*N
chosen = []
heapq.heapify(chosen)

tot = 0
num_kinds = 0
ans = [0]*(N+1)

for i in range(K):
    t, d = TD[i]
    tot += d
    Tcount[t] += 1
    if Tcount[t] == 1: 
        num_kinds += 1
    if Tcount[t] > 1:
        heapq.heappush(chosen, d)

ans[num_kinds] = tot + num_kinds**2

for i in range(K, N):
    t, d = TD[i]
    if Tcount[t] > 0: continue
    Tcount[t] += 1
    num_kinds += 1
    tot = tot - heapq.heappop(chosen) + d
    ans[num_kinds] = tot + num_kinds**2
    if len(chosen) == 0:
        break

print(max(ans))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
TD = []
for _ in range(N):
    t, d = map(int, input().split())
    TD.append((t-1, d))

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

import heapq

Tcount = [0]*N
chosen = []
heapq.heapify(chosen)

tot = 0
num_kinds = 0
ans = [0]*N

for i in range(K):
    t, d = TD[i]
    tot += d
    Tcount[t] += 1
    if Tcount[t] == 1: 
        num_kinds += 1
    if Tcount[t] > 1:
        heapq.heappush(chosen, d)

ans[num_kinds] = tot + num_kinds**2

for i in range(K, N):
    t, d = TD[i]
    if Tcount[t] > 0: continue
    Tcount[t] += 1
    num_kinds += 1
    tot = tot - heapq.heappop(chosen) + d
    ans[num_kinds] = tot + num_kinds**2
    if len(chosen) == 0:
        break

print(max(ans))

Submission status: Runtime Error

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
    if l<0:break
    continue
  if TOP[data[r][0]] != 0:
    r += 1
    if r>=N:break
    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: Runtime Error

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)

Submission status: Runtime Error

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: Runtime Error

Solution code:
N,K = map(int,input().split())
TD = [tuple(map(int,input().split())) for i in range(N)]
TD.sort(key=lambda x:x[1])

import heapq
hq = []
heapq.heapify(hq)
selected = set()
ans = 0
for _ in range(K):
    t,d = TD.pop()
    if t in selected:
        heapq.heappush(hq, d)
    else:
        selected.add(t)
    ans += d
ans += len(selected)**2

tmp = ans
while TD:
    t,d = TD.pop()
    if t in selected: continue
    selected.add(t)
    tmp -= heapq.heappop(hq)
    tmp += d
    tmp += len(selected)*2 - 1
    ans = max(ans ,tmp)
print(ans)

Submission status: Runtime Error

Solution code:
from heapq import *

n, k = map(int,input().split())
T = [[] for i in range(n)]
for i in range(n):
    t,d = map(int,input().split())
    T[t-1].append(d)
Top = []
for i in range(n):
    if T[i]:
        T[i].sort(reverse=True)
        Top.append((T[i][0],i))
Top.sort(reverse = True)

various = len(Top)


ans = 0
s = 0
sushi = []
heapify(sushi)
l = 0

for j,(t,i) in enumerate(Top[:k]):
    s += t
    if l + j + 1 > k:
        s -= heappop(sushi)
        l -= 1
    for new in T[i][1:]:
        if l + j + 1 >= k:
            s += new
            s -= heappushpop(sushi,new)
            if sushi[0] >= new:break
        else:
            s += new
            l += 1
            heappush(sushi, new)
    ans = max(ans, s + (j+1)**2)
print(ans)

Submission status: Runtime Error

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

I = []
for _ in range(N):
  t,d = map(int, input().split())
  I.append((d,t))
  
I.sort(reverse=True)

# heap_can_delete
h_d = []
# heap_can_add
h_a = []

wk = 0
kind = 0
# 出現したことのあるネタ
ap = set()
for i in range(K):
  d,t = I[i]
  wk += d
  if t not in ap:
    ap.add(t)
    kind += 1
  else:
    heapq.heappush(h_d, (d,t))

for i in range(K,N):
  d,t = I[i]
  if t not in ap:
    ap.add(t)
    # maxをpopさせるために、価値を-1かけておく
    heapq.heappush(h_a, ((-1) * d, t))
  
ans = wk + kind ** 2
while len(h_d) > 0:
  d1,t1 = heapq.heappop(h_d)
  d2,t2 = heapq.heappop(h_a)
  d2 = (-1) * d2
  wk = wk - d1 + d2
  kind += 1
  ans = max(ans, wk + kind ** 2)
  
print(ans)

Submission status: Runtime Error

Solution code:
from heapq import *
import sys
I=sys.stdin.readlines()
N,K=map(int,I[0].split())
t,d=0,0
X=[]
for i in range(N):
  t,d=map(int,I[i+1].split())
  X.append((d,t))
X.sort()
Q=[]
heapify(Q)
C=[0]*N
V=0
S=0
for i in range(K):
  heappush(Q,X[-1])
  if C[X[-1][1]-1]==0:
    V+=1
  S+=X[-1][0]
  C[X[-1][1]-1]+=1
  del X[-1]
def f(q,c,s,x,v):
  y=0
  r=s+v*v
  for i in range(len(q)):
    y=heappop(q)
    if c[y[1]-1]>=2:
      s-=y[0]
      c[y[1]-1]-=1
      break
    if len(q)==0:
      return r
  for i in range(len(x)):
    if c[x[-1][1]-1]==0:
      heappush(q,x[-1])
      s+=x[-1][0]
      c[x[-1][1]-1]=1
      v+=1
      del x[-1]
      return max(r,f(q,c,s,x,v))
    del x[-1]
  return r
print(f(Q,C,S,X,V))

Submission status: Runtime Error

Solution code:
N,K = map(int,input().split())
dic = {}
A = []
for _ in range(N):
    t,d = map(int,input().split())
    A.append([t,d])
A.sort(key = lambda x:x[1], reverse = True)
val = 0
set_ = set()
duplicated = collections.deque([])
for t,d in A[:K]:
    val += d
    if t in set_:
        duplicated.append(d)
    else:
        set_.add(t)
now = val + len(set_)**2

nx = collections.deque(A[K:])
ans = now
c = len(set_)
while len(set_) < K and nx:
    i,j = nx.popleft()
    if i in set_:
        continue
    d = duplicated.pop()
    diff = j-d + 2*c+1
    set_.add(i)
    now += diff
    ans =max(ans,now)
    c += 1
print(ans)

Submission status: Runtime Error

Solution code:
import collections
N,K = map(int,input().split())
dic = {}
A = []
for _ in range(N):
    t,d = map(int,input().split())
    A.append([t,d])
A.sort(key = lambda x:x[1],reverse = True)
val = 0
set_= set()
D = collections.deque([])
for t,d in A[:K]:
    val += d
    if t in set_:
        D.append(d)
    else:
        set_.add(t)
now = val + len(set_)**2
nx = collections.deque(A[K:])
ans = now
c = len(set_)
while len(set_)<K and nx:
    i,j = nx.popleft()
    if i in set_:
        continue
    d = D.pop()
    diff = j-d+2*c+1
    set_.add(j)
    now += diff
    ans = max(ans,now)
    c +=1
print(ans)

Submission status: Runtime Error

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):
		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: Runtime Error

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] # 削除候補
    w = -heapq.heappop(hq)
    if v < w:
        cur += w - v
    else:
        heapq.heappush(hq, -w)
    ans = max(ans, cur + i*i)
print(ans)

Submission status: Runtime Error

Solution code:
import heapq
N, K = map(int,input().split())
L1 = []
L2 = []
S = set()
val = 0

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

L1.sort(reverse = True)

for i in range(K):
    val += L1[i][0]
    if L1[i][1] in S:
        heapq.heappush(L2, -L1[i][0])
    else:
        S.add(L1[i][1])

ans = val + len(S) ** 2

for i in range(K, N):
    if not L1[i][1] in S:
        S.add(L1[i][1])
        val += heapq.heappop(L2)
        val += L1[i][0]
        ans = max(ans, val + len(S) ** 2)

print(ans)