Problem:


Score : 400 points

Problem Statement

There are N pieces of sushi. Each piece has two parameters: "kind of topping" t_i and "deliciousness" d_i.
You are choosing K among these N pieces to eat.
Your "satisfaction" here will be calculated as follows:

The satisfaction is the sum of the "base total deliciousness" and the "variety bonus".

The base total deliciousness is the sum of the deliciousness of the pieces you eat.

The variety bonus is x*x, where x is the number of different kinds of toppings of the pieces you eat.

You want to have as much satisfaction as possible.
Find this maximum satisfaction.

Constraints

1 \leq K \leq N  \leq 10^5

1 \leq t_i \leq N

1 \leq d_i \leq 10^9

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N K
t_1 d_1
t_2 d_2
.
.
.
t_N d_N

Output

Print the maximum satisfaction that you can obtain.

Sample Input 1

5 3
1 9
1 7
2 6
2 5
3 1

Sample Output 1

26

If you eat Sushi 1,2 and 3:

The base total deliciousness is 9+7+6=22.

The variety bonus is 2*2=4.

Thus, your satisfaction will be 26, which is optimal.

Sample Input 2

7 4
1 1
2 1
3 1
4 6
4 5
4 5
4 5

Sample Output 2

25

It is optimal to eat Sushi 1,2,3 and 4.

Sample Input 3

6 5
5 1000000000
2 990000000
3 980000000
6 970000000
6 960000000
4 950000000

Sample Output 3

4900000016

Note that the output may not fit into a 32-bit integer type.



Sample input:
5 3
1 9
1 7
2 6
2 5
3 1


Sample output:
26

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys

def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dp = [[[0,0,set()] for j in range(K+1)] for i in range(N+1) ]

    for i in range(N):
        for k in range(1,K+1):
            di = i + 1
            
            p_dp = dp[di-1][k]
            c_dp = dp[di][k]

            p_k  = len(p_dp[2])
            c_k  = p_k + 1 if not t[i] in p_dp[2] else p_k

            p_p  = p_dp[0]
            c_p  = p_dp[0] - p_k**2 + c_k**2 + d[i]

            if p_p < c_p and p_dp[1] < k:
                dp[di][k][0] = c_p
                dp[di][k][1] = dp[di-1][k][1] + 1
                dp[di][k][2] = p_dp[2].copy()
                dp[di][k][2].add(t[i])
            else:
                dp[di][k][0] = dp[di-1][k][0]
                dp[di][k][1] = dp[di-1][k][1] 
                dp[di][k][2] = dp[di-1][k][2].copy() 

    print(dp[N][K][0])
    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: Wrong Answer

Solution code:
from collections import Counter
from operator import itemgetter

N, K = list(map(int, input().split()))
td = sorted([list(map(int, input().split())) for i in range(N)], key=itemgetter(1))

group = td[N-K:]
count = Counter()
td = td[:N-K]

total = 0
for t, d in group:
    count[t] += 1
    total += d

ans = total + len(count)**2

td.reverse()

for t, d in td:
    changed = False
    for t_, d_ in group:
        if count[t_] > 1:
            changed = True
            count[t] += 1
            total = total - d_ + d
            ans = max(ans, total+len(count)**2)
            group.remove([t_, d_])
            group.append([t, d])
            group = sorted(group, key=itemgetter(1))
            break
    if not changed:
        break
print(ans)

Submission status: Wrong Answer

Solution code:
from collections import Counter
from operator import itemgetter

N, K = list(map(int, input().split()))
td = sorted([list(map(int, input().split())) for i in range(N)], key=itemgetter(1))
group = sorted(td[-K:])
td = td[:N-K]

count = Counter([group[i][0] for i in range(K)])
t_sum = sum([group[i][1] for i in range(K)])
ans = t_sum + len(count)**2

for t, d in td:
    flag = False
    for t_, d_ in group:
        if count[t_] > 1:
            flag = True
            count[t_] -= 1
            t_sum -= d_
            group.remove([t_, d_])
            group.append([t, d])
            count[t] += 1
            t_sum += d
            group = sorted(group, key=itemgetter(1))
            ans = max(ans, t_sum+len(count)**2)
            break
    if not flag:
        break
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq
n,k = map(int,input().split())
td = [[0,0] for i in range(n)]
for i in range(n):
    td[i][0],td[i][1] = map(int,input().split())
td.sort(reverse = True,key = lambda x:x[1])
heap = [];sushitype = set();oisisa = 0
for i in range(k):
    oisisa += td[i][1]
    heapq.heappush(heap,td[i][1])
    sushitype.add(str(td[i][0]))
for i in range(k+1,n):
    if str(td[i][0]) in sushitype:
        continue
    n = heapq.heappop(heap)
    if len(sushitype) ** 2 + oisisa < (len(sushitype) + 1) ** 2 + oisisa - n + td[i][1]:
        sushitype.add(td[i][0])
        heapq.heappush(heap,td[i][1])
        oisisa = oisisa - n + td[i][1]
    else:
        heapq.heappush(heap,n)

print(sum(heap) + len(sushitype) ** 2)
"""
print(td)
print(heap)
"""

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
"""
D - Various Sushi
https://atcoder.jp/contests/abc116/tasks/abc116_d

"""
import sys

from collections import Counter

def solve(N, K, sushi):
    sushi.sort(key=lambda x: x[1], reverse=True)
    pick = sushi[:K]            #  おいしさポイントの高い方からK個の寿司を選択
    pick_type = set([t for t, d in pick]) # 選択済の種類

    cnt = Counter()             #  各種類の寿司が何個ずつあるか覚えておく
    for t, d in pick:
        cnt[t] += 1

    rem = dict()                #  未選択の種類の寿司のうち、一番おいしさポイントが高い物
    for t, d in sushi:
        if t not in pick_type:
            rem[t] = max(rem.get(t, 0), d)
    res = sorted([[k, v] for k, v in rem.items()], key=lambda x: x[1]) # おいしさで昇順にソート(一番おいしい物が末尾)

    manzoku = sum([d for t, d in pick])
    ans = [manzoku + len(pick_type)**2] #  現在の選択でのスコア、ここから種類を1つずつ増加させていきスコアを計算する
    while len(pick_type) < K and pick and res:
        t, d = res.pop()
        pick_type.add(t)
        tt, dd = pick.pop()
        if cnt[tt] > 1:
            cnt[tt] -= 1
            manzoku = manzoku - dd + d
            ans.append(manzoku+ len(pick_type)**2)
    return max(ans)


def main(args):
    N, K = map(int, input().split())
    sushi = [[int(i) for i in input().split()] for _ in range(N)]
    ans = solve(N, K, sushi)
    print(ans)


if __name__ == '__main__':
    main(sys.argv[1:])

Submission status: Wrong Answer

Solution code:
N, K = [int(n) for n in input().split()]

def calc(items):
    # print(items)
    t, p = zip(*items)
    t_len = len(set(t))
    tv = t_len ** 2
    pv = sum(p)
    return tv + pv

def main():
    
    items = [tuple(map(lambda s: int(s), input().split())) for _ in range(K)]
    ans = calc(items)

    for _ in range(N - K):
        curr_items = items[:]
        item = tuple(map(lambda s: int(s), input().split()))
        
        for i in range(K):
            old_item = curr_items.pop(i)
            curr_items.insert(i, item)
            new_ans = calc(curr_items)
            if new_ans > ans:
                items = curr_items
                ans = new_ans
            else:
                curr_items.pop(i)
                curr_items.insert(i, old_item)

    print(ans)

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
N, K = [int(n) for n in input().split()]

def calc(items):
    # print(items)
    t, p = zip(*items)
    t_len = len(set(t))
    tv = t_len ** 2
    pv = sum(p)
    return tv + pv

def main():
    
    items = [tuple(map(lambda s: int(s), input().split())) for _ in range(K)]
    ans = calc(items)

    for _ in range(N - K):
        item = tuple(map(lambda s: int(s), input().split()))
        
        for i in range(K):
            curr_items = items[:]
            curr_items.pop(i)
            curr_items.append(item)
            new_ans = calc(curr_items)
            if new_ans > ans:
                items = curr_items
                ans = new_ans

    print(ans)

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
total = len(base_t)
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  if len(base_t)==total:
    break
  
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    elm = remains.pop(tg_idx)
    base.pop(rm_idx)
    base.append(elm)
    remains = remains[tg_idx:]
    base_t.add(elm[0])
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int,input().split(' '))
sushis = sorted([tuple(map(int, input().split(' '))) for _ in range(n)], key=lambda x:x[1], reverse=True)
total_t = len(set( t for t, d in sushis ))
base = sushis[:k]
base_t = set( t for t, d in base )
remains = sushis[k:]

ans = sum( d for t, d in base ) + len(set( t for t, d in base))**2

while True:
  rm_idx = k
  for i in range(len(base)-1, -1, -1):
    if len(set( t for t, d in base )) == len(set( t for t, d in base[:i] + base[i+1:] )):
      rm_idx = i
      break
  
  tg_idx = -1
  for i in range(len(remains)):
    if remains[i][0] not in base_t:
      tg_idx = i
      break
  
  if rm_idx==k or tg_idx==-1:
    break
  else:
    base[rm_idx] = remains.pop(tg_idx)
    ans = max(ans , sum( d for t, d in base ) + len(set( t for t, d in base))**2)

print(ans)

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys
from math import *
from itertools import *
from collections import *
from functools import *
from heapq import *
try:
    from math import gcd
except Exception:
    from fractions import gcd


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    dt = sorted(zip(d, t), reverse=True)

    q0 = []

    acc = 0
    ts = 0
    p = [-1] * (N+1)
    for dd, tt in dt[:K]:
        acc += dd
        v = min(p[tt], dd)
        p[tt] = max(p[tt], dd)
        if 0 < v:
            heappush(q0, (v, tt))
        else:
            ts += 1
    
    q1 = []
    for dd, tt in dt[K:]:
        heappush(q1, (-dd, tt))

    mmm = acc + ts * ts
#    while 0 < len(q1):
#        d1, t1 = heappop(q1)
#        if 0 < p[t1]:
#            continue
#        acc -= d1
#        p[t1] = -d1
#        ts += 1
#
#        d0, _ = heappop(q0)
#        acc -= d0
#        mmm = max(mmm, acc + ts * ts)

    return mmm


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))
    result = solve(N, K, t, d)
    print(result)

if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())

t = []
td = []
for i in range(n):
  ti, di = map(int, input().split())
  t.append(ti)
  td.append([ti, di])
  
td = sorted(td, key = lambda x: x[1], reverse=True)

ans = 0
for i in range(1, min(k, len(set(t)))+1):
  p = i * i
  num = 0
  x = []
  for j in range(n):
    if len(x) < i and (td[j][0] not in x):
      num += 1
      x.append(td[j][0])
      p += td[j][1]
    elif len(x) <= i and (k - i > num - len(x)):
      num += 1
      p += td[j][1]
    if num >= k or len(x) >= i:
      break
  ans = max(ans, p)
print(ans)

Submission status: Wrong Answer

Solution code:
import sys
sys.setrecursionlimit(10**7)
INF = 10 ** 18
MOD = 10 ** 9 + 7
from functools import partial, reduce
from operator import mul
prod = partial(reduce, mul)
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 input()
from collections import Counter

# fibo = [1, 1]
# for i in range(2, 9):
#     fibo.append(fibo[i - 1] + fibo[i - 2])

def main():
    N, K = LI()
    TD = Counter()
    for _ in range(N):
        t, d = LI()
        TD[t] = max(TD[t], d)
    ds = list(TD.values())
    ds.sort(reverse=True)
    tmp = ds[0] * K
    ans = tmp + 1
    for i in range(1, min(K, len(ds))):
        tmp = tmp + ds[i] - ds[0]
        ans = max(ans, tmp + ((i+1)**2))
    return ans

print(main())

Submission status: Wrong Answer

Solution code:
# 漸化式的に解く
# 一般性を失わない範囲をちゃんと考える。
# "置き換え可能"というのをどの部分で持たせられるかを考える
#　今回の場合は、リスト作りの際に作ったので十分

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = set()
change = [] #変更できるもの(2個以上あるもの)



cnt = 0
for t,d in pairs[:K]: #K個上から順番に取ってみる
    if t in kinds:
        change.append(d) 
    kinds.add(t)
    cnt += d

n = len(kinds)
cnt += len(kinds)
ans = cnt

"""
K+1以降
Kまでに大きい順に加えたもので、2個以上あるものから順に置き換えていく。
1個のものはそのままで良い。
また、K+1以降で、2回同じものを追加する可能性はない。
なぜなら新規の追加ではないので、その場合はchangeに入ってある美味しさdiよりも
K+1以降のものは必ず低い美味しさの値になっているからである。
"""
for t,d in pairs[K:]:
    # 交換できるものがなくなったら終わり
    if len(change) == 0:
        break
    # 同じものは交換しない
    if t in kinds:
        continue
    d_min = change.pop()
    cnt = cnt - (d_min + n**2)  + (d + (n+1)**2)
    kinds.add(t)
    n += 1
    ans = max(cnt,ans)
# 集計
print(ans)

Submission status: Wrong Answer

Solution code:
# 漸化式的に解く
# 一般性を失わない範囲をちゃんと考える
# 2種類以上あるものは、同じ種類のものとして差し支えない

from collections import defaultdict
N, K = map(int,input().split())
# dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]
# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = [0] * (N+1)

oisii = []
for t,d in pairs[:K]: #K個上から順番に取ってみる
    kinds[t] += 1
    oisii.append([t,d])
point = sum([d for t,d in oisii]) + sum([1 for kind in kinds if kind!=0])**2
ans = point

for t,d in pairs[K:]:
    """
    置き換えようとしていたけど
    その必要はなく、
    ほしいのはmax値だけなので、配列そのものを更新し続ける必要はない
    """
    # if kinds[t] == 1:
    #     continue
    # else:
    #     index = K
    #     while True:
    #         t_min, d_min = oisii[index-1]
    #         # 2個以上ある種類のなかで最小値を求める
    #         if kinds[t_min] >= 2:
    #             break
    #         index -= 1
    #         if index == 0:
    #             break
    # if index == 0:
    #     break
    if len(oisii) == 0:
        break
    if kinds[t] >= 1:
        continue
    # とりあえず最小だったoisiiを置き換えていく
    t_min, d_min = oisii.pop()
    n = sum([1 for kind in kinds if kind!=0])
    # とりあえず種類を追加
    kinds[t] += 1
    kinds[t_min] -= 1
    point = point - (d_min + n**2)  + (d + (n+1)**2)
    # print(point)
    ans = max(point,ans)
# 集計
print(ans)

Submission status: Wrong Answer

Solution code:
# 漸化式的に解く
# 一般性を失わない範囲をちゃんと考える
# 2種類以上あるものは、同じ種類のものとして差し支えない

from collections import defaultdict
N, K = map(int,input().split())
dd = defaultdict(list)
pairs = [list(map(int,input().split())) for _ in range(N)]

# 美味しい順にソート
pairs = sorted(pairs,reverse=True, key=lambda x: x[1])
kinds = [0] * (N+1)

cnt = 0
oisii = []
for t,d in pairs[:K]: #K個上から順番に取ってみる
    kinds[t] += 1
    cnt += d
    # print(t,d)
    oisii.append([t,d])

point = sum([d for t,d in oisii]) + sum([1 for kind in kinds if kind!=0])**2
ans = point
for t,d in pairs[K:]:
    if kinds[t] == 1:
        continue
    else:
        index = K
        while True:
            t_min, d_min = oisii[index-1]
            # 2個以上ある種類のなかで最小値を求める
            if kinds[t_min] >= 2:
                break
            index -= 1
            if index == 0:
                break
        if index == 0:
            break
        n = sum([1 for kind in kinds if kind!=0])
        # とりあえず種類を追加
        oisii = [o for i,o in enumerate(oisii) if i != index]
        oisii.append([t,d])
        kinds[t] += 1
        kinds[t_min] -= 1
        point = point - (d_min + n**2)  + (d + (n+1)**2)
        ans = max(point,ans)
# 集計
print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
td = []
for i in range(N):
    _td = list(map(int, input().split()))
    td.append(_td)
td.sort(key=lambda x:x[1], reverse=True)

# 美味しさ順
base_var = set([_td[0] for _td in td[:K]])
score_d = sum([_td[1] for _td in td[:K]]) + len(base_var)**2

score_list = [score_d]

candidate_list = []
for _td in td[K:]:
    if _td[0] not in base_var:
        candidate_list.append(_td)

current = td[:K]

for candidate in candidate_list:
    current_var = set([_td[0] for _td in current])
    for i in range(len(current))[::-1]:
        _current = current[::]
        _current.pop(i)
        if len(current_var) != len(set([_td[0] for _td in _current])):
            continue
        _current.insert(0, candidate)
        var = set([_td[0] for _td in _current])
        score = sum([_td[1] for _td in _current]) + len(var) ** 2
        score_list.append(score)
        current = _current
        break

print(max(score_list))

Submission status: Wrong Answer

Solution code:
N,K = map(int, input().split())
sushi = [[int(i) for i in input().split()] for _ in range(N)]

dp = [[0 for _ in range(K+2)] for _ in range(N+2)]
sushi_set = [[set() for _ in range(K+2)] for _ in range(N+2)] #対応するdpの場所でどの種類の寿司を保持しているか

for i in range(1,N+1):
  for j in range(1,K+1):
    tmp = set([sushi[i-1][0]])
    tmp_sushi_set = sushi_set[i-1][j-1] | tmp
    cmp = dp[i-1][j-1] - len(sushi_set[i-1][j-1])**2 + sushi[i-1][1] + len(tmp_sushi_set)**2

    if dp[i-1][j] > cmp:
      dp[i][j] = dp[i-1][j]
    else:
      dp[i][j] = cmp
      sushi_set[i][j] = tmp_sushi_set

    
print(dp[N][K])

Submission status: Wrong Answer

Solution code:
import sys
from collections import defaultdict

input = sys.stdin.readline

N, K = map(int, input().split())
sushi = defaultdict(list)
for i in range(N):
    t, d = map(int, input().split())
    sushi[t].append(d)

max_sushi, remain_sushi = [], []
for v in sushi.values():
    v = sorted(v, reverse=True)
    max_sushi.append(v[0])
    if len(v) > 1:
        remain_sushi += v[1:]
max_sushi = sorted(max_sushi, reverse=True)

ans = 0
for x in range(1, K + 1):
    a = max_sushi[:x]
    b = remain_sushi[:K - x]
    p = sum(a) + sum(b) + x * x
    ans = max(ans, p)
print(ans)

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
L = []
for _ in range(N):
    t = tuple(map(int,input().split()))
    t = (t[1],t[0])
    L.append(t)
L = sorted(L,reverse=True)
h = L[:K]
L = L[K:]
using = [0 for _ in range(100000+10)]
determined = []
s = set() 
total = 0
for x in h:
    using[x[1]] += 1
    s.add(x[1])
    total += x[0]
cand = []
n = len(s)
cand.append((total,n))
import heapq as hp
hp.heapify(h)
j = 0
for _ in range(len(h)):
    tmp = hp.heappop(h)
    if using[tmp[1]] == 1:
        determined.append(tmp)
    else:
        for i in range(j,len(L)):
            if L[i][1] in s:
                pass
            else:
                using[tmp[1]] -= 1
                j = i+1
                s.add(x[1])
                total -= tmp[0]
                total += L[i][0]
                n += 1
                cand.append((total,n))
ans = []
for x in cand:
    ans.append(x[1]**2+x[0])
print(max(ans))

Submission status: Wrong Answer

Solution code:
N,K=map(int, input().split())
sushi=sorted([[int(x) for x in input().split()] for _ in range(N)], key=lambda x: x[1], reverse=True)

sty = set()
ims = []
mus = []

for t,d in sushi[:K]:
	if t in sty:
		mus.append(d)
	else:
		sty.add(t)
		ims.append(d)

prs = []

for t,d in sushi[K:]:
	if t not in sty:
		sty.add(t)
		prs.append(d)

ans = len(ims) * len(ims) + sum(ims) + sum(mus)

while len(mus) > 0 and len(prs) > 0:
	mus.pop()
	ims.append(prs.pop())
	ans = max(ans, len(ims) * len(ims) + sum(ims) + sum(mus))

print(ans)

Submission status: Wrong Answer

Solution code:
N,K=map(int, input().split())
sushi=sorted([[int(x) for x in input().split()] for _ in range(N)], key=lambda x: x[1], reverse=True)

sty = set()
ims = []
mus = []

for t,d in sushi[:K]:
	if t in sty:
		mus.append(d)
	else:
		sty.add(t)
		ims.append(d)

prs = []

for t,d in sushi[K:]:
	if t not in sty:
		sty.add(t)
		prs.append(d)

ans = len(ims) * len(ims) + sum(ims + mus)

while len(mus) * len(prs) > 0:
	mus.pop()
	ims.append(prs.pop())
	ans = max(ans, len(ims) * len(ims) + sum(ims + mus))

print(ans)

Submission status: Wrong Answer

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

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

td.sort(reverse=True, key=lambda x: x[1])

sumd0 = 0
sumd1 = 0
sd0 = [0]
sd1 = [0]
added_t = []
for i in range(N):
    t, d = td[i]
    if t in added_t:
        sumd0 += d
        sd0.append(sumd0)
    else:
        sumd1 += d
        sd1.append(sumd1)
        added_t.append(t)
    prev = t

S = [0] * (K+1)
for i in range(K+1):
    S[i] = sd1[min(i, len(sd1)-1)] + sd0[min(K-i, len(sd0)-1)] + (i * i)

ans = max(S)
print(ans)

Submission status: Wrong Answer

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

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

td.sort(reverse=True, key=lambda x: x[1])

sumd0 = 0
sumd1 = 0
sd0 = [0]
sd1 = [0]
prev = -1
for i in range(N):
    t, d = td[i]
    if t == prev:
        sumd0 += d
        sd0.append(sumd0)
    else:
        sumd1 += d
        sd1.append(sumd1)
    prev = t

S = [0] * (K+1)
for i in range(K+1):
    S[i] = sd1[min(i, len(sd1)-1)] + sd0[min(K-i, len(sd0)-1)] + (i * i)

ans = max(S)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq

N, K = map(int, input().split())
heap = []
types = [0 for _ in range(N)]

class Sushi():
    def __init__(self, neta, aji):
        self.neta = neta - 1
        self.aji = aji

    def __str__(self):
        return "Sushi{0}: {1}".format(self.neta, self.aji)

    def reverse(self):
        self.aji = -self.aji

    def __lt__(self, other):
        return self.aji < other.aji

for _ in range(N):
    t, d = map(int, input().split())
    heap.append(Sushi(t, -d))
heapq.heapify(heap)

variety = 0
choices = []
for _ in range(K):
    sushi = heapq.heappop(heap)
    sushi.reverse()
    neta = sushi.neta
    if types[neta] == 0:
        variety += 1
    types[neta] += 1
    choices.append(sushi)
heapq.heapify(choices)

basic = 0
for c in choices:
    basic += c.aji
satisfaction = basic + variety * variety
top_satisfaction = satisfaction
while True:
    drop_sushi = None
    while len(choices) > 0:
        choiced_sushi = heapq.heappop(choices)
        if types[choiced_sushi.neta] > 1:
            drop_sushi = choiced_sushi
            types[choiced_sushi.neta] -= 1
            break
    add_sushi = None
    while len(heap) > 0:
        new_sushi = heapq.heappop(heap)
        new_sushi.reverse()
        if types[new_sushi.neta] == 0:
            add_sushi = new_sushi
            break
    if drop_sushi is None or add_sushi is None:
        break
    else:
        variety += 1
        basic += add_sushi.aji - drop_sushi.aji
        satisfaction = basic + variety * variety
        if satisfaction > top_satisfaction:
            top_satisfaction = satisfaction
print(top_satisfaction)

Submission status: Wrong Answer

Solution code:
nk = input().split()
n = int(nk[0])
k = int(nk[1])

dtlist = []
for i in range(n):
    td = input().split()
    t = int(td[0])
    d = int(td[1])
    dtlist.append((d, t))
dtlist.sort(reverse=True)


t_to_maxidx = {}
for idx, dt in enumerate(dtlist):
    d, t = dt
    if t not in t_to_maxidx:
        t_to_maxidx[t] = idx


selected = []
count = 0
tcount = {}

for t, idx in t_to_maxidx.items():
    if count >= k:
        break
    selected.append(dtlist[idx])
    count += 1
    tcount[t] = 1 if t not in tcount else tcount[t] + 1
    dtlist[idx] = None

new_dtlist = []
for dt in dtlist:
    if dt is not None:
        new_dtlist.append(dt)
dtlist = new_dtlist


if k - count > 0:
    for i in range(k - count):
        dt = dtlist.pop(0)
        selected.append(dt)
        d, t = dt
        tcount[t] = 1 if t not in tcount else tcount[t] + 1

selected.sort()

score = 0
for d, t in selected:
    score += d

len_tcount = len(tcount)
score += len_tcount * len_tcount

for i in range(0, 2**32):
    next_score = score
    try:
        d_new, t_new = dtlist.pop(0)
        d_old, t_old = selected.pop(0)
    except IndexError:
        print(score)
        break
    next_score += d_new
    next_score -= d_old

    tcount[t_old] -= 1
    if tcount[t_old] <= 1:
        next_score -= len_tcount * len_tcount
        len_tcount -= 1
        next_score += len_tcount * len_tcount
        del tcount[t_old]

    if t_new not in tcount:
        next_score -= len_tcount * len_tcount
        len_tcount += 1
        next_score += len_tcount * len_tcount
        tcount[t_new] = 1
    else:
        tcount[t_new] += 1

    if next_score <= score:
        print(score)
        break
    score = next_score

Submission status: Wrong Answer

Solution code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = sorted([(t, d) for t, d in zip(td[::2], td[1::2])], key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    count = [0] * n
    for _c in choice:
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_remove, i_append = k - 1, k
    while len(variety) < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        while count[td[i_remove][0] - 1] == 1:
            i_remove -= 1
        variety.add(td[i_append][0])
        count[td[i_remove][0] - 1] -= 1
        ans.append(ans[-1] - td[i_remove][1] + td[i_append][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Wrong Answer

Solution code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = [(t, d) for t, d in zip(td[::2], td[1::2])]
    td = sorted(td, key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    count = [0] * n
    for _c in choice:
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_remove, i_append = k - 1, k
    while len(variety) < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        while count[td[i_remove][0] - 1] == 1:
            i_remove -= 1
        variety.add(td[i_append][0])
        ans.append(ans[-1] - td[i_remove][1] + td[i_append][1] + 2 * len(variety) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Wrong Answer

Solution code:
def solve(string):
    n, k, *td = map(int, string.split())
    td = [(t, d) for t, d in zip(td[::2], td[1::2])]
    td = sorted(td, key=lambda x: x[1], reverse=True)
    max_variety = len(set([c[0] for c in td]))
    choice = td[:k]
    count = [0] * n
    for _c in choice:
        count[_c[0] - 1] += 1
    variety = set([c[0] for c in choice])
    ans = [len(variety)**2 + sum([c[1] for c in choice])]
    i_remove, i_append = k - 1, k
    while len(variety) + len(ans) - 1 < min(max_variety, k):
        while td[i_append][0] in variety:
            i_append += 1
        while count[td[i_remove][0] - 1] == 1:
            i_remove -= 1
        ans.append(ans[-1] - td[i_remove][1] + td[i_append][1] + 2 * (len(variety) + len(ans)) - 1)
    return str(max(ans))


if __name__ == '__main__':
    n, m = map(int, input().split())
    print(solve('{} {}\n'.format(n, m) + '\n'.join([input() for _ in range(n)])))

Submission status: Wrong Answer

Solution code:
n,k = [int(i) for i in input().split()]
td = [[int(i) for i in input().split()] for _ in range(n)]

###dが大きい順にソート
td.sort(key = lambda x: x[0],reverse=True)
td.sort(key = lambda x: x[1],reverse=True)

###なるべくtの種類を多いなかで、dを最大化する
###まず各tでdが最大のものを1つずつ選ぶ
eated_kind = []
eated_kind_num = 0
eated = []
not_eated = [i for i in td]
for i in range(k):
  t = not_eated[0][0]
  eated_kind.append(t)
  eated_kind_num += 1
  eated.append(not_eated[0])
  not_eated = [i for i in not_eated if i[0] != t]
    
  ###tがk種類存在したら、残りのtを扱うことはない
###  ###またdも最大化されているため、これ以降のtdを使うことはない
  if eated_kind_num == k:
    break
  if len(not_eated) == 0:
    break

not_eated = [i for i in td if i not in eated]

###tがk種類存在しなかったら、k個になるまでtに関係なくdを最大化する
if eated_kind_num < k:
  eated += not_eated[:k-eated_kind_num]
  not_eated = not_eated[k-eated_kind_num:]
  #print(eated)
  eated.sort(key = lambda x: x[1],reverse=True)
  
###dの最大化を目指すが、tの種類が減る場合どちらのほうが良いか判断する必要がある
for t,d in not_eated:
  eated_last_t, eated_last_d = eated[-1]
  if d < eated_last_d:
    break
  t_point_sub = (eated_kind_num-1)**2 - eated_kind_num**2
  d_point_sub = d - eated_last_d
  if t_point_sub + d_point_sub > 0:
    eated = [[t,d]] + eated[:-1]
    eated_kind_num -= 1
  else:
    break
    
#print(eated)

t_point = eated_kind_num**2
d_point = sum([i[1] for i in eated])
point = t_point + d_point
print(point)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter
from collections import Counter
import heapq

N, K = [int(v) for v in input().split()]
sushi = [tuple(int(v) for v in input().split()) for i in range(N)]

sushi = sorted(sushi, key=itemgetter(1), reverse=True)

ate_list = sushi[:K]

kind_counter = Counter([s[0] for s in ate_list])
kind_cnt = len(list(kind_counter))

max_pt = sum([s[1] for s in ate_list]) + kind_cnt ** 2

if K >= N:
    print(max_pt)
    exit()

print(ate_list)

waiting_sushi = []
wait_kind = []
for s in sushi[K:]:
    if kind_counter[s[0]]:  # すでに食べてる種類
        continue

    if not s[0] in wait_kind:
        waiting_sushi.append(s)
        wait_kind.append(s[0])

# 食べた種類を増やしていくフェーズ
idx = K - 1

print(waiting_sushi)

curr_pt = max_pt
for s in waiting_sushi:
    target = None
    for i in range(idx, -1, -1):
        # 1個しかないものは交換しない
        if kind_counter[ate_list[i][0]] > 1:
            target = ate_list[i]
            idx = i
            break

    if not target:
        break

    # 入れ替えして比較
    val = curr_pt - target[1] + s[1] + ((kind_cnt + 1) ** 2) - (kind_cnt ** 2)
    curr_pt = val

    max_pt = max(curr_pt, max_pt)
    kind_counter[target[0]] -= 1

    kind_cnt += 1
    idx -= 1

# print(kind_cnt)
print(max_pt)