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.

Solution (Python):
from heapq import heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()
hp = []
for t, d in TD:
    flg = 0
    if t not in muptiple:
        muptiple.add(t)
        flg = 1
    heappush(hp, (d*-1, flg))

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

total = vari ** 2 + point

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

print(total)

Solution (Python):
import operator
import collections

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

sushies.sort(key=operator.itemgetter(1), reverse=True)

dishes = sushies[:k]
c = collections.Counter(t for t, d in dishes)
a = sum(d for t, d in dishes)
answer = a + len(c) ** 2
x = k - 1
for t, d in sushies[k:]:
    if t in c:
        continue

    while x >= 0:
        if c[dishes[x][0]] > 1:
            break
        x -= 1

    if x == -1:
        break


    a = a + d - dishes[x][1]
    c[t] += 1
    c[dishes[x][0]] -= 1
    del dishes[x]
    dishes.append((t, d))

    answer = max(answer, a + len(c) ** 2)

print(answer)

Solution (Python):
import heapq

n,k=map(int,input().split())
arr=[list(map(int,input().split())) for _ in range(n)]
arr=sorted(arr,reverse=True,key=lambda x:x[1])
q=[]
v=set()
s=0
for t,d in arr[:k]:
  s+=d
  if t in v:
    heapq.heappush(q,d)
  else:
    v.add(t)
s+=len(v)**2
ans=s
for t,d in arr[k:]:
  if t not in v and len(q)!=0:
    z=heapq.heappop(q)
    s+=d+2*len(v)+1-z
    v.add(t)
    ans=max(ans,s)
print(ans)

Solution (Python):
def main():

    """ convenient functions
    
    
    
    
    
    """

    """convenient decorator
    
    
        
        
        
        
        
        
        
        
        
        
    """
    
    
    import sys
    sys.setrecursionlimit(10**7)
    from itertools import accumulate, combinations, permutations, product 
    
    from math import factorial, ceil, floor, sqrt
    def factorize(n):
        """return the factors of the Arg and count of each factor
        
        Args:
            n (long): number to be resolved into factors
        
        Returns:
            list of tuples: factorize(220) returns [(2, 2), (5, 1), (11, 1)]
        """
        fct = []  
        b, e = 2, 0  
        while b * b <= n:
            while n % b == 0:
                n = n // b
                e = e + 1
            if e > 0:
                fct.append((b, e))
            b, e = b + 1, 0
        if n > 1:
            fct.append((n, 1))
        return fct
    def combinations_count(n, r):
        """Return the number of selecting r pieces of items from n kinds of items.
        
        Args:
            n (long): number
            r (long): number
        
        Raises:
            Exception: not defined when n or r is negative
        
        Returns:
            long: number
        """
        
        if n < 0 or r < 0:
            raise Exception('combinations_count(n, r) not defined when n or r is negative')
        if n - r < r: r = n - r
        if r < 0: return 0
        if r == 0: return 1
        if r == 1: return n
        numerator = [n - r + k + 1 for k in range(r)]
        denominator = [k + 1 for k in range(r)]
        for p in range(2,r+1):
            pivot = denominator[p - 1]
            if pivot > 1:
                offset = (n - r) % p
                for k in range(p-1,r,p):
                    numerator[k - offset] /= pivot
                    denominator[k] /= pivot
        result = 1
        for k in range(r):
            if numerator[k] > 1:
                result *= int(numerator[k])
        return result
    def combinations_with_replacement_count(n, r):
        """Return the number of selecting r pieces of items from n kinds of items allowing individual elements to be repeated more than once.
        
        Args:
            n (long): number
            r (long): number
        
        Raises:
            Exception: not defined when n or r is negative
        
        Returns:
            long: number
        """
        if n < 0 or r < 0:
            raise Exception('combinations_with_replacement_count(n, r) not defined when n or r is negative')
        elif n == 0:
            return 1
        else:
            return combinations_count(n + r - 1, r)
    from bisect import bisect_left, bisect_right
    from collections import deque, Counter, defaultdict 
    from heapq import heapify, heappop, heappush, heappushpop, heapreplace,nlargest,nsmallest 
    from copy import deepcopy, copy 
    import operator
    from operator import itemgetter 
    
    
    from functools import reduce, lru_cache
    def chmin(x, y):
        """change minimum
        if x > y, x = y and return (x, True).
        convenient when solving problems of dp[i]
        
        Args:
            x (long): current minimum value
            y (long): potential minimum value
        
        Returns:
            (x, bool): (x, True) when updated, else (x, False)
        """
        if x > y:
            x = y
            return (x, True)
        else:
            return (x, False)
    def chmax(x, y):
        """change maximum
        if x < y, x = y and return (x, True).
        convenient when solving problems of dp[i]
        
        Args:
            x (long): current maximum value
            y (long): potential maximum value
        
        Returns:
            (x, bool): (x, True) when updated, else (x, False)
        """
        if x < y:
            x = y
            return (x, True)
        else:
            return (x, False)

    from fractions import gcd 
    def gcds(numbers):
        return reduce(gcd, numbers)
    def lcm(x, y):
        return (x * y) // gcd(x, y)
    def lcms(numbers):
        return reduce(lcm, numbers, 1)

    
    
    INF = 10 ** 18
    MOD = 10 ** 9 + 7
    modpow = lambda a, n, p = MOD: pow(a, n, p) 
    def modinv(a, p = MOD):
        
        
        return modpow(a, p-2, p)
    def modinv_list(n, p = MOD):
        if n <= 1:
            return [0,1][:n+1]
        else:
            inv_t = [0,1]
            for i in range(2, n+1):
                inv_t += [inv_t[p % i] * (p - int(p / i)) % p]
            return inv_t
    def modfactorial_list(n, p = MOD):
        if n == 0:
            return [1]
        else:
            l = [0] * (n+1)
            tmp = 1
            for i in range(1, n+1):
                tmp = tmp * i % p
                l[i] = tmp
            return l
    def modcomb(n, k, fac_list = [], p = MOD):
        
        
        from math import factorial
        if n < 0 or k < 0 or n < k: return 0
        if n == 0 or k == 0: return 1
        if len(fac_list) <= n:
            a = factorial(n) % p
            b = factorial(k) % p
            c = factorial(n-k) % p
        else:
            a = fac_list[n]
            b = fac_list[k]
            c = fac_list[n-k]
        return (a * modpow(b, p-2, p) * modpow(c, p-2, p)) % p
    def modadd(a, b, p = MOD):
        return (a + b) % MOD
    def modsub(a, b, p = MOD):
        return (a - b) % p
    def modmul(a, b, p = MOD):
        return ((a % p) * (b % p)) % p
    def moddiv(a, b, p = MOD):
        return modmul(a, modpow(b, p-2, p))

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

    """ initialize variables and set inputs
    
        
        
    
        
        
        
        
        
        
        
    
        
        
        
        
    
        
        
    """

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

    """ how to treat input
    
    
    
    
    
    
    
    """

    if __name__ != '__main__':
        
        sys.stdin = open('sample.txt')
        
    
    

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

    """memo: how to use defaultdict of list
    
    Dic = defaultdict(list)
    
    Dic[x].append(y)
    
    for k, v in Dic.items():
    """
    
    """memo: how to solve binary problems
    
    >>> a = 5
    >>> bin_str_a = format(a, '
    >>> print(bin_str_a)
    0b0101 
    """
    
    """memo: how to solve the problem
    creating simple test/answer
    greed
    simple dp
    graph
    """

if __name__ == '__main__':
    main()

Solution (Python):
from heapq import heappush, heappop

N, K = map(int, input().split())
sushi = [[] for _ in range(N+1)]
for _ in range(N):
    t, d = map(int, input().split())
    sushi[t].append(d)
for li in sushi:
    li.sort()

first = [(li[-1], i) for i, li in enumerate(sushi) if li]
first.sort(reverse = True)

kind = 0
v_first = 0
second = []
second_sum = 0
second_cnt = 0

ans = 0
for v, i in first:
    kind += 1
    if kind > K:
        break
    v_first += v
    for d in sushi[i][:-1]:
        heappush(second, d)
        second_sum += d
        second_cnt += 1
    while second_cnt > (K-kind):
        d = heappop(second)
        second_cnt -= 1
        second_sum -= d
    value = kind*kind + v_first + second_sum
    if ans < value:
        ans = value

print(ans)

Solution (Python):
import sys
input = sys.stdin.readline

n, k = map(int, input().split())
dt = []
for _ in range(n):
    a,b = map(int, input().split())
    dt.append([b,a])
dt.sort(reverse=True)


t = [0] * (n+1)
from heapq import *

h = []
h2 = []
f = [0] * (n+1)
g = [0] * (n+1)
c = 0
tf = 0
for i in range(k):
    tmp = dt[i]
    if t[tmp[1]] == 0:
        t[tmp[1]] += 1
        c += 1
    else:
        heappush(h, tmp[0])
    tf += tmp[0]
f[c] = tf

for i in range(k, n):
    tmp = dt[i]
    if t[tmp[1]] == 0:
        t[tmp[1]] += 1
        heappush(h2, -tmp[0])

while len(h2) > 0 and len(h) > 0:
    tout = heappop(h)
    tin = -heappop(h2)
    c += 1
    tf = tf - tout + tin
    f[c] = tf

for i in range(n+1):
    ff = f[i]
    if ff != 0:
        g[i] = ff + i*i
print(max(g))

Solution (Python):
from heapq import heappop, heappush
from collections import defaultdict

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


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


idx = K
while pq:
    
    d_p, t_p = heappop(pq)
    if chosen[t_p] <= 1:
        continue
    chosen[t_p] -= 1
    score -= d_p

    
    while idx < N and chosen[X[idx][0]] != 0:
        idx += 1

    if idx == N:
        break
        
    cnt += 1
    chosen[X[idx][0]] += 1
    score += X[idx][1]
    ans = max(ans, score + cnt ** 2)
    
print(ans)

Solution (Python):



import sys
import math
import bisect
from heapq import heappush, heappop
from collections import defaultdict
from itertools import accumulate
from collections import Counter
from collections import deque
from operator import itemgetter
mod = 10**9 + 7

n,k = map(int,input().split())
td = [list(map(int,input().split())) for _ in range(n)]
td.sort(key = itemgetter(1),reverse = True)
f = [0]*n
ans = 0
kind = 0
x = []
for i in range(k):
    t,d = td[i]
    ans += d
    if not f[t-1]:
        kind += 1
    else:
        x.append(d)
    f[t-1] += 1
ans += kind**2
x.sort()
x = deque(x)
y = ans
for i in range(k,n):
    t,d = td[i]
    if not x:
        break
    if f[t-1]:
        continue
    y += d - x.popleft() + 2*kind + 1
    f[t-1] = 1
    kind += 1
    ans = max(ans,y)
print(ans)

Solution (Python):
N,K = map(int,input().split())
TD = [list(map(int,input().split())) for _ in [0]*N]
TD.sort(key=lambda x:(-x[1],x[0]))
use = TD[:K]
var = set()
lock = []
unlock = []
for t,d in use:
    if t in var:
        unlock.append(d)
    else:
        var.add(t)
        lock.append(d)
unlock.sort(reverse=True)
Dscore = sum(lock) + sum(unlock)
v = len(var)
Vscore = v*v
ans = Dscore + Vscore
unused = TD[K:][::-1]
while unlock and unused:
    t,d = unused.pop()
    if t in var: continue
    trash = unlock.pop()
    var.add(t)
    lock.append(d)
    v += 1
    Dscore += d - trash
    Vscore = v*v
    ans = max(ans,Dscore + Vscore)
print(ans)

Solution (Python):

import heapq
import sys
input = sys.stdin.readline

def main():
    num_sushi, chose_num = map(int, input().split())
    data = [list(map(int, input().split())) for i in range(num_sushi)]

    data.sort(key=lambda x: -x[1])
    delicious_heap = []
    count_d, count_n = 0, 0
    neta_data = [0 for i in range(num_sushi + 1)]

    for i in range(chose_num):
        neta, oishisa = data[i]

        if neta_data[neta]:
            heapq.heappush(delicious_heap, oishisa)
        else:
            neta_data[neta] = 1
            count_n += 1

        count_d += oishisa

    ans = count_d + count_n ** 2

    for i in range(chose_num, num_sushi):
        if len(delicious_heap) == 0:
            break
        neta, oishisa = data[i]

        if neta_data[neta]:
            continue
        neta_data[neta] = 1
        count_n += 1
        pop_oishisa = heapq.heappop(delicious_heap)
        count_d += oishisa - pop_oishisa

        ans = max(ans, count_d + count_n ** 2)

    print(ans)


if __name__ == '__main__':
    main()

Solution (Python):

import heapq


def main():
    num_sushi, chose_num = map(int, input().split())
    data = [list(map(int, input().split())) for i in range(num_sushi)]

    data.sort(key=lambda x: -x[1])
    delicious_heap = []
    count_d, count_n = 0, 0
    neta_data = [0 for i in range(num_sushi + 1)]

    for i in range(chose_num):
        neta, oishisa = data[i]

        if neta_data[neta]:
            heapq.heappush(delicious_heap, oishisa)
        else:
            neta_data[neta] = 1
            count_n += 1

        count_d += oishisa

    ans = count_d + count_n ** 2

    for i in range(chose_num, num_sushi):
        if len(delicious_heap) == 0:
            break
        neta, oishisa = data[i]

        if neta_data[neta]:
            continue
        neta_data[neta] = 1
        count_n += 1
        pop_oishisa = heapq.heappop(delicious_heap)
        count_d += oishisa - pop_oishisa

        ans = max(ans, count_d + count_n ** 2)

    print(ans)


if __name__ == '__main__':
    main()

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

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

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

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

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

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

if k >= 0:
    res = max(res, kind ** 2 + base + rest)

print(res)

Solution (Python):
from heapq import heappush, heappop

N, K = map(int, input().split())
t, d = (
    zip(*(map(int, input().split()) for _ in range(N))) if N else
    ((), ())
)




S = sorted(zip(d, t), reverse=True)
q = []
v = set()
s = 0
for D, T in S[:K]:
    s += D
    if T in v:
        heappush(q, D)
    else:
        v.add(T)
s += len(v)**2
opt = s
for D, T in S[K:]:
    if T not in v and q:
        z = heappop(q)
        s += D + 2 * len(v) + 1 - z
        v.add(T)
        opt = max(opt, s)
ans = opt

print(ans)

Solution (Python):
import sys
sys.setrecursionlimit(10**7)
def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x) - 1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def II(): return int(sys.stdin.readline())
def SI(): return sys.stdin.readline().strip()
INF = 10 ** 18
MOD = 10 ** 9 + 7

def score(td):
  t_list = [v[0] for v in td]
  d_list = [v[1] for v in td]
  ut = len(set(t_list))
  score_ut = ut**2
  score_d = sum(d_list)
  return score_ut + score_d
  
def solve():
  N, K = LI()
  td = []
  for i in range(N):
    t, d = LI()
    td.append([t, d])

  from operator import itemgetter
  td.sort(key=itemgetter(1), reverse=True)

  
  cand = td[:K]
  
  
  import collections
  c = collections.Counter([v[0] for v in cand])
  m = len(c.keys())
  
  
  rest_top1 = {}
  for v in td:
    if v[0] not in c:
      if v[0] not in rest_top1:
      	rest_top1[v[0]] = v[1]
      else:
        if rest_top1[v[0]] < v[1]:
          rest_top1[v[0]] = v[1]
          
  
  rest_o = sorted([v for k, v in rest_top1.items()], reverse=True)
  current_score = score(cand)
  scores = [current_score]
  if len(rest_o) == 0:
    print(current_score)
    return
  
  nget = 0
  for i in range(1, K):
    target = cand[-i]
    if c[target[0]] == 1:
      continue 
    
    
    
    dscore_neta = (m+nget+1)**2 - (m+nget)**2 
    
    dscore_d = -target[1] + rest_o[nget]
    current_score += dscore_neta
    current_score += dscore_d
    
    scores.append(current_score)
    
    nget += 1
    c[target[0]] -= 1
    
    
    if nget == len(rest_o):
      break
  print(max(scores))
solve()

Solution (Python):
import heapq

N, K = map(int, input().split())
S = []
for i in range(N):
    t, d = map(int, input().split())
    heapq.heappush(S, (-d, t))

sum_d = 0
types = set()
surplus = []
for i in range(K):
    _d, t = heapq.heappop(S)
    sum_d += -_d
    if t in types:
        heapq.heappush(surplus, -_d)
    else:
        types.add(t)
len_types = len(types)
ans = sum_d + len_types * len_types

while S and surplus:
    _d, t = heapq.heappop(S)
    if t in types:
        continue
    types.add(t)
    removed_d = heapq.heappop(surplus)
    sum_d += -_d - removed_d
    len_types += 1
    ans = max(ans, sum_d + len_types * len_types)
print(ans)

Solution (Python):
import heapq

N, K = map(int, input().split())
S = []
for i in range(N):
    t, d = map(int, input().split())
    heapq.heappush(S, (-d, t))

sum_d = 0
types = set()
surplus = []
for i in range(K):
    _d, t = heapq.heappop(S)
    sum_d += -_d
    if t in types:
        heapq.heappush(surplus, -_d)
    else:
        types.add(t)
len_types = len(types)
ans = sum_d + len_types * len_types

while S and surplus:
    _d, t = heapq.heappop(S)
    if t in types:
        continue
    types.add(t)
    removed_d = heapq.heappop(surplus)
    sum_d += -_d - removed_d
    len_types += 1
    ans = max(ans, sum_d + len_types * len_types)
print(ans)

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

    first = set()
    second = []
    point = 0
    for s in sushi[:k]:
        if s[0] not in first:
            first.add(s[0])
        else:
            second.append(s)
        point += s[1]
    point += len(first) ** 2
    ans = point

    for s in sushi[k:]:
        if len(second) == 0:
            break
        if s[0] not in first:
            point = point + s[1] - second.pop()[1] + \
                (len(first) + 1) ** 2 - len(first) ** 2
            first.add(s[0])
            ans = max(ans, point)
    print(ans)


main()

Solution (Python):
import sys
def input():
    return sys.stdin.readline()[:-1]
from heapq import *
from collections import defaultdict
N, K = map(int, input().split())
d_t = [list(map(int, input().split()))[::-1] for i in range(N)]
d_t.sort(reverse = 1)
f, c = 0, 0
dic = defaultdict(lambda : 0)
take = []
for d, t in d_t[:K]:
    f += d
    if dic[t] == 0:
        c += 1
    dic[t] += 1
    heappush(take, [d, t])
ans = f + c * c
rest = [[-d, t] for d, t in d_t[K:]]
heapify(rest)
while take:
    d, t = heappop(take)
    if dic[t] == 1: continue
    dic[t] -= 1
    f -= d
    while rest:
        rd, rt = heappop(rest)
        if dic[rt] == 0:
            dic[rt] = 1
            c += 1
            f -= rd
            ans = max(ans, f + c * c)
            break
print(ans)

Solution (Python):
from collections import defaultdict
import sys

input = sys.stdin.readline

n, k = map(int, input().split())
e = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    t -= 1
    e[t].append(d)

tops = []
rests = []
for l in e.values():
    for i, d in enumerate(sorted(l, reverse=True)):
        if i == 0:
            tops.append(d)
        else:
            rests.append(d)

tops.sort(reverse=True)
c_tops = [0]
for t in tops:
    c_tops.append(c_tops[-1] + t)

rests.sort(reverse=True)
c_rests = [0]
for t in rests:
    c_rests.append(c_rests[-1] + t)

ans = 0
for i in range(1, min(k, len(tops)) + 1):
    if k - i > len(rests): continue
    ans = max(ans, c_tops[i] + c_rests[k - i] + pow(i, 2))
print(ans)

Solution (Python):
from heapq import heappush, heappop
from operator import itemgetter
import sys

input = sys.stdin.readline

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

e.sort(key=itemgetter(0), reverse=True)

b = e[:k]  

used = [False] * n
hq = []

ans = 0  
tn = 0  
for d, t in b:
    if not used[t]:
        used[t] = True
        tn += 1
    else:
        heappush(hq, (d, t))
    ans += d
ans += pow(tn, 2)  

c = e[k:]  





s = ans
cur = 0
while hq:
    d, t = heappop(hq)
    s -= d
    changed = False
    while cur < n - k:
        cd, ct = c[cur]
        cur += 1
        if not used[ct]:
            used[ct] = True
            changed = True
            s += cd
            s += tn * 2 + 1
            tn += 1
            if s > ans:
                ans = s
            break
    if not changed:
        break
print(ans)

Solution (Python):
import heapq

N, K = map(int, input().split())
sushi = []
for i in range(N):
    t, d = map(int, input().split())
    heapq.heappush(sushi, (-d, t))

total = 0
types = set()
remain = []
for i in range(K):
    d, t = heapq.heappop(sushi)
    total -= d
    if t not in types:
        types.add(t)
    else:
        
        heapq.heappush(remain, -d)

types_len = len(types)
ans = total + types_len ** 2

while sushi and remain:
    d, t = heapq.heappop(sushi)
    if t not in types:
        types.add(t)
        remain_d = heapq.heappop(remain)
        total = total - d - remain_d
        types_len += 1
        ans = max(ans, total + types_len ** 2)

print(ans)

Solution (Python):
from operator import itemgetter
n,k = map(int,input().split())
li=[]
s=[]
t=[]
g=0
ans=[]
retu=[0]*n
p=0
q=0
for _ in range(n):
    a,b=map(int,input().split())
    li.append([a,b])
li=sorted(li,key=itemgetter(1),reverse=True)
for i in range(k):
    g+=li[i][1]
    if retu[li[i][0]-1]==0:
        retu[li[i][0]-1]+=1
        p+=1
    else:
        t.append(li[i][1])
        q+=1
g+=p**2
ans.append(g)
for i in range(k,n):
    if q==0:
        break
    if retu[li[i][0]-1]==0:
        g=g-t[q-1]+li[i][1]+2*p+1
        p+=1
        q-=1
        ans.append(g)
        retu[li[i][0]-1]+=1
print(max(ans))

Solution (Python):
N , K = map(int,input().split())

sushi = []

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

sushi.sort(key=lambda A : A[1], reverse=True)


nowSushi = []
sumTaste = 0
numKind = 0
take = [0] * N  

for i in range(K) :
    t, d = sushi[i]
    nowSushi.append((t,d))
    sumTaste += d
    if take[t-1] == 0 :
        numKind += 1
    take[t-1] += 1

addSushi = []  
for i in range(K,N) :
    t, d = sushi[i]
    if take[t-1] > 0 :  
        continue
    addSushi.append((t,d))
    take[t-1] += 1

ans = sumTaste + numKind**2
L = len(addSushi)
changeCount = 0

while nowSushi :
    removeSushiT, removeSushiD = nowSushi.pop()
    if take[removeSushiT-1] <= 1 :
        continue
    if changeCount >= L :
        break

    addSushiT, addSushiD = addSushi[changeCount]
    changeCount += 1

    sumTaste = sumTaste - removeSushiD + addSushiD
    numKind = numKind + 1

    if (sumTaste + numKind**2) > ans :
        ans = sumTaste + numKind**2

print(ans)

Solution (Python):
import sys
input = sys.stdin.readline
'''
allinputs = iter(input().splitlines())
input = lambda : next(allinputs)



def caluculate_score():
	return num_kind * num_kind + taste
	
	
N, K = map(int, input().split())
t, d = [0] * N, [0] * N
table = [[] for _ in range(N)]
eat_kind = [0] * N
num_kind = 0

for i in range(N):
	t[i], d[i] = map(int, input().split())
	table[i] = (d[i], t[i] - 1)
	
table.sort()

eat = []
taste = 0

for _ in range(K):
	tmp = table.pop()
	taste += tmp[0]
	eat_kind[tmp[1]] += 1
	if eat_kind[tmp[1]] == 1:
		num_kind += 1
	eat.append(tmp)
	
ans = caluculate_score()
ind_eat = K - 1
ind_table = N - K - 1

while ind_eat >= 0 and ind_table >= 0:
	while ind_eat >= 0 and eat_kind[eat[ind_eat][1]] <= 1:
		ind_eat -= 1
		
	while ind_table >= 0 and eat_kind[table[ind_table][1]] >= 1:
		ind_table -= 1
		
	if not(ind_eat >= 0 and ind_table >= 0):
		break

	eat_kind[eat[ind_eat][1]] -= 1
	eat_kind[table[ind_table][1]] += 1
	num_kind += 1
	taste += table[ind_table][0] - eat[ind_eat][0]
	eat[ind_eat], table[ind_table] = table[ind_table], eat[ind_eat]
	score = caluculate_score()
	if score > ans:
		ans = score
		
	ind_eat -= 1
	ind_table -= 1

print(ans)

Solution (Python):



n, k = list(map(int, input().split()))

sushi = [list(map(int, input().split())) + [None] for _ in range(n)]
sushi.sort(key = lambda s: s[1] , reverse = True) 


neta_seen = set()
for i in range(n):
	if sushi[i][0] in neta_seen:
		sushi[i][2] = False
	else:
		sushi[i][2] = True
		neta_seen.add(sushi[i][0])




n_kind = 0
for i in range(k):
	if sushi[i][2]:
		n_kind += 1

ans = sum([s[1] for s in sushi[:k]]) + n_kind ** 2


now = ans
adopt = k
trash = k-1

while(True):

	while(adopt < n and sushi[adopt][2] == False):
		adopt += 1
	while(trash >= 0 and sushi[trash][2] == True):
		trash -= 1
	if not(adopt < n and trash >= 0):
		break

	now = now + sushi[adopt][1] - sushi[trash][1]
	now = now - n_kind**2 + (n_kind+1) ** 2



	n_kind += 1
	adopt += 1
	trash -= 1

	if now > ans:
		ans = now

print(ans)

Solution (Python):
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
total = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    total -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, -d)

length = len(group)
ans = length**2 + total

while td and td_:
    d, t = heapq.heappop(td)

    if t not in group:
        group.add(t)
        d_ = heapq.heappop(td_)
        total = total - d_ - d
        length += 1
        ans = max(ans, length**2 + total)
print(ans)

Solution (Python):
import heapq

N, K = list(map(int, input().split()))
td = []
td_ = []
group = set()
total = 0

for i in range(N):
    t, d = list(map(int, input().split()))
    heapq.heappush(td, (-d, t))

for i in range(K):
    d, t = heapq.heappop(td)
    total -= d
    if t not in group:
        group.add(t)
    else:
        heapq.heappush(td_, -d)

ans = len(group)**2 + total

while td and td_:
    d, t = heapq.heappop(td)

    if t not in group:
        group.add(t)
        d_ = heapq.heappop(td_)
        total = total - d_ - d
        ans = max(ans, len(group)**2 + total)
print(ans)