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: Accepted

Solution code:
# ABC 116 D

import heapq

def resolve():
    n, k = map(int, input().split())
    td = []

    # td に score が高い順に格納する
    for _ in range(n):
        t, d = map(int, input().split())
        heapq.heappush(td, (-d, t))

    score = 0
    types = set()
    discard = []

    for _ in range(k):
        d, t = heapq.heappop(td)  # 注意 d は符号が逆になってる
        score -= d
        if t not in types:
            types.add(t)
        else:
            # その type で1番じゃないものをスコアの小さい順に並べる -> 後で捨てるリスト
            heapq.heappush(discard, (-d, t))

    ans = score + len(types)**2

    while td and discard:
        d, t = heapq.heappop(td)  # 残り物からスコア高い順に取り出す; 注意 このリストのスコアは符号反転してる

        if t in types:
            # 選んでも種類が増えないやつはパス
            continue

        types.add(t)
        score -= d

        # 捨てても種類が減らないリストの中で最もスコアが低いやつを捨てる
        _d, _t = heapq.heappop(discard)  # 注意 このリストのスコアは符号反転していない
        score -= _d

        ans = max(ans, score + len(types)**2)

    print(ans)

if __name__ == "__main__":
    resolve()

Submission status: Accepted

Solution code:
from operator import itemgetter

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

TD.sort(key=itemgetter(1), reverse=True)

remain = []
removable = []
used = [False for _ in range(N+1)]
val = 0
S = 0
for i, (t, d) in enumerate(TD):
    if i < K:
        S += d
        if not used[t]:
            val += 1
        else:
            removable.append(d)
        used[t] = True
    else:
        if not used[t]:
            remain.append(d)
            used[t] = True

S += val**2
removable.sort()

delta = 0
ans = S
for k, r in enumerate(remain):
    if k >= len(removable):
        break
    val += 1
    delta += r - removable[k] + 2*val - 1
    ans = max(ans, S+delta)
print(ans)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
import sys, math, fractions, itertools, heapq

def solve(N: int, k: int, t: "List[int]", d: "List[int]"):
    arr = sorted(zip(t, d), key=lambda x: x[1], reverse=True)
    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 # arr[:k]での点数
    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)
    return


# Generated by 1.1.4 https://github.com/kyuridenamida/atcoder-tools  (tips: You use the default template now. You can remove this line by using your custom template)
def main():
    def iterate_tokens():
        for line in sys.stdin:
            for word in line.split():
                yield word
    tokens = iterate_tokens()
    N = int(next(tokens))  # type: int
    K = int(next(tokens))  # type: int
    t = [int()] * (N)  # type: "List[int]" 
    d = [int()] * (N)  # type: "List[int]" 
    for i in range(N):
        t[i] = int(next(tokens))
        d[i] = int(next(tokens))
    solve(N, K, t, d)

if __name__ == '__main__':
    main()

Submission status: Accepted

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

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

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

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

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
N, K = map(int, input().split())
TD = [list(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda x:x[1], reverse=True)
muptiple = set()
for i in range(len(TD)):
    flg = 0
    if TD[i][0] not in muptiple:
        muptiple.add(TD[i][0])
        flg = 1
    TD[i][0] = flg
TD.reverse()
ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = TD.pop()
    vari += v[0]
    point += v[1]
    if not v[0]:
        heappush(ret, v[1])

total = vari ** 2 + point

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

print(total)

Submission status: Accepted

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

ret = []
point = 0
vari = 0
total = 0
for i in range(K):
    v = heappop(hp)
    t = v[0]*-1
    vari += v[1]
    point += t
    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)

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:
def main():

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

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

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

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

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

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

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

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

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

    N, K = R()
    TD = [Rtuple() for _ in range(N)]
    TD.sort(key=itemgetter(1), reverse=True)
    kind_sushi = set()
    num_kind = 0
    num_eaten = 0
    skip_candidate = []
    ans = 0
    tmp = 0
    for t, d in TD:
        if num_eaten == K:
            if t in kind_sushi or not skip_candidate:
                continue
            num_kind += 1
            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)
    
    # end of main
    # ----------------------------------

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

if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:
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)
# print(dt)

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))

Submission status: Accepted

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

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

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

# Step 2: Replace chosen items
idx = K
while pq:
    # Step 2.1
    d_p, t_p = heappop(pq)
    if chosen[t_p] <= 1:
        continue
    chosen[t_p] -= 1
    score -= d_p

    # Step 2.2
    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)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
#ABC116 D

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)

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:

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()

Submission status: Accepted

Solution code:

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()

Submission status: Accepted

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

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

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

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

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

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

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

print(res)

Submission status: Accepted

Solution code:
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
    ((), ())
)

# おいしさについて貪欲に寿司を食べたときの満足ポイントを求める
# その後、種類ボーナスを1個ずつ増やしたときの満足ポイントをそれぞれ求め、
# それらの最大値を求める
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)

Submission status: Accepted

Solution code:
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)

  # おいしさtopK
  cand = td[:K]
  
  # cand のネタ保持数dict
  import collections
  c = collections.Counter([v[0] for v in cand])
  m = len(c.keys())
  
  # 持ってないネタのtop1リスト
  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
  # おいしく無い順に持ってないもののtop1と入れ替え
  nget = 0
  for i in range(1, K):
    target = cand[-i]
    if c[target[0]] == 1:
      continue # 一つしか持ってないなら入れ替えない
    
    # スコアの差分を作る
    # netaボーナス
    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()

Submission status: Accepted

Solution code:
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)

Submission status: Accepted

Solution code:
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()

Submission status: Accepted

Solution code:
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)