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:
import heapq


def func(n, k, td):
    neta = [0 for _ in range(n + 1)]
    td.sort()
    td.reverse()
    _heapq = []
    _max = 0
    _sum = 0
    types = 0
    for i in range(k):
        _sum += td[i][0]
        if neta[td[i][1]] == 0:
            types += 1
        if neta[td[i][1]] > 0:  # 被っているものだけ後で削除する可能性あり
            heapq.heappush(_heapq, td[i])
        neta[td[i][1]] += 1
    _max = _sum + types**2
    # print("td:{} heapq:{}".format(td, _heapq))
    for i in range(k, n):
        # neta[td[i][1]] != 0なら,
        # td[i] < td[j](j<i)なのでtd[i]は無視
        if len(_heapq) > 0 and neta[td[i][1]] == 0:  # 種類が増えるとき
            tmp = heapq.heappop(_heapq)
            _sum -= tmp[0]
            _sum += td[i][0]
            neta[tmp[1]] -= 1
            types += 1
        # print("td:{} heapq:{}".format(td, _heapq))
        # print("_sum:{} typed:{}".format(_sum, types))
        _max = max(_max, _sum + types**2)
    return str(_max)


n, k = [int(i) for i in input().split()]
td = [[0, 0] for _ in range(n)]
for i in range(n):
    td[i][1], td[i][0] = [int(i) for i in input().split()]
# print("print(func({},{},{}))".format(n, k, td))
print(func(n, k, td))
"""
print("26 vs "+func(5,3,[[9, 1], [7, 1], [6, 2], [5, 2], [1, 3]]))
print("25 vs "+func(7,4,[[1, 1], [1, 2], [1, 3], [6, 4], [5, 4], [5, 4], [5, 4]]))
print("4900000016 vs "+func(6,5,[[1000000000, 5], [990000000, 2], [980000000, 3], [970000000, 6], [960000000, 6], [950000000, 4]]))
"""

Submission status: Wrong Answer

Solution code:
def manzoku(s):
    return sum([d for (t, d) in s]) + len(set([t for (t, d) in s]))**2

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))
#print(menu)
    
#Sort with value d
sorted(menu, key = lambda x:x[1], reverse=True)
#print(menu)

menu_ = menu[:K]
#print(menu_)
#x = len(set([t for (t, d) in menu_]))
_menu = menu[K:]
#print(_menu)

#print(menu_)
ans = manzoku(menu_)
flg = True
for idx in range(K, N):
    for (t, d) in reversed(menu_):
        if [i for (i, v) in menu_].count(t) > 1:
            menu_.remove((t, d))
            #print(menu_)
            for (ii, vv) in _menu:
                if ii not in [i for (i, v) in menu_]:
                    menu_.append((ii, vv))
                    #print(menu_)
                    _menu.remove((ii, vv))
                    #print(_menu)
                    flg = False
                    break
        if flg == False:
            break

    p = manzoku(menu_)
    #print(ans, p)
    ans = max(ans, p)

print(ans)

Submission status: Wrong Answer

Solution code:
def manzoku(s):
    return sum([d for (t, d) in s]) + len(set([t for (t, d) in s]))**2

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))
    
#Sort with value d
sorted(menu, key = lambda x:x[1], reverse=True)
#print(menu)

menu_ = menu[:K]
x = len(set([t for (t, d) in menu_]))
_menu = menu[K:]

def sushi():
    for (t, d) in reversed(menu_):
        if [i for (i, v) in menu_].count(t) != 1:
            menu_.remove((t, d))
            #print(menu_)
            for (ii, vv) in _menu:
                if ii != t:
                    menu_.append((ii, vv))
                    #print(menu_)
                    _menu.remove((ii, vv))
                    #print(_menu)
                    return False
    return True        

#print(menu_)
g = []
for idx in range(x, N+1):
    g.append(manzoku(menu_))

    if sushi():
        break

print(max(g))

Submission status: Wrong Answer

Solution code:
def manzoku(s):
    return sum([d for (t, d) in s]) + len(set([t for (t, d) in s]))**2

N, K = map(int, input().split())
menu = []

for i in range(N):
    t, d = map(int, input().split())
    menu.append((t, d))
    
#Sort with value d
sorted(menu, key = lambda x:x[1], reverse=True)
#print(menu)

menu_ = menu[:K]
_menu = menu[K:]

def sushi():
    for (t, d) in reversed(menu_):
        if [i for (i, v) in menu_].count(t) != 1:
            menu_.remove((t, d))
            #print(menu_)
            for (ii, vv) in _menu:
                if ii != t:
                    menu_.append((ii, vv))
                    #print(menu_)
                    _menu.remove((ii, vv))
                    #print(_menu)
                    return False
    return True        

#print(menu_)
g = []
for idx in range(N):
    g.append(manzoku(menu_))

    if sushi():
        break

print(max(g))

Submission status: Wrong Answer

Solution code:
import heapq
N, K = map(int, input().split())
dt = []
for i in range(N):
	t, d = map(int, input().split())
	dt.append((d, t))
dt.sort()
dt.reverse()
dic = {}
can_delete = []
dp = 0
x = 0
for i in range(K):
	d, t = dt[i]
	dp += d
	if t in dic:
		heapq.heappush(can_delete, d)
	else:
		dic[t] = d
		x += 1
ans = dp + x * x
v = ans
for i in range(K, N):
	if len(can_delete) == 0:
		break
	d, t = dt[i]
	if t in dic:
		continue
	x += 1
	v = v - heapq.heappop(can_delete) + d + 2 * x - 1
	if v > ans:
		ans = v
print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
h = []
for i in range(n):
    t, d = map(int, input().split())
    h.append((t,d))
h.sort(key=lambda x: -x[1])

taste = 0
typen = 0
typecnt = {}
bad = []


for i in range(k):
    typecnt[h[i][0]] = typecnt.get(h[i][0], 0) + 1
    if(typecnt[h[i][0]] == 1):
        typen += 1
    else:
        bad.append(h[i][1])
    taste += h[i][1]


ans = taste + typen**2

for i in range(k, n):
    if(len(bad) == 0):
        break
    typen += 1
    taste -= bad.pop()
    taste += h[i][1]
    ans = max(ans, taste + typen**2)

print(ans)

Submission status: Wrong Answer

Solution code:
import collections

n, k = map(int, input().split())
h = []
for i in range(n):
    a, b = map(int, input().split())
    h.append((a,b))
h.sort(key=lambda x: -x[1])
fav=h[0:k]
tb={}
for t,d in sorted(h[k:], key=lambda x: (-x[0], -x[1])):
    tb[t] = tb.get(t, []) + [d]
alt = []
for t,ds in tb.items():
    alt.append((t,max(ds)))
alt.sort(key=lambda x: -x[1])

# print("h: ",h)
# print("fav: ",fav)
# print("alt: ",alt)


var = len(set([x[0] for x in fav]))

bestscore = var**2 + sum([x[1] for x in fav])

while(len(alt) > 0):
    c = collections.Counter([x[0] for x in fav])
    temp = sorted(filter(lambda x: c[x[0]] > 1, fav), key=lambda x: -x[1])
    if(len(temp) == 0):
        break
    while(len(alt) > 0):
        t1, d1 = alt.pop(0)
        if(c[t1] == 0):
            break
    t2, d2 = temp[-1]


    fav.remove((t2,d2))
    fav.append((t1, d1))
    var += 1

    bestscore = max(bestscore, var**2 + sum([x[1] for x in fav]))

    # print("fav: ",fav)
    # print("alt: ",alt)

print(bestscore)

Submission status: Wrong Answer

Solution code:
import collections

n, k = map(int, input().split())
h = []
for i in range(n):
    a, b = map(int, input().split())
    h.append((a,b))
h.sort(key=lambda x: -x[1])
fav=h[0:k]
tb={}
for t,d in sorted(h[k:], key=lambda x: (-x[0], -x[1])):
    tb[t] = tb.get(t, []) + [d]
alt = []
for t,ds in tb.items():
    alt.append((t,max(ds)))
alt.sort(key=lambda x: -x[1])

# print("h: ",h)
# print("fav: ",fav)
# print("alt: ",alt)
while(len(alt) > 0):
    c = collections.Counter([x[0] for x in fav])
    temp = sorted(filter(lambda x: c[x[0]] > 1, fav), key=lambda x: -x[1])
    if(len(temp) == 0):
        break
    while(len(alt) > 0):
        t1, d1 = alt.pop(0)
        if(c[t1] == 0):
            break
    t2, d2 = temp[-1]

    ts = set([x[0] for x in fav])
    if d1 - d2 + (len(ts)+1)**2 - len(ts)**2 > 0:
        fav.remove((t2,d2))
        fav.append((t, d))
        # print("fav: ", fav)
    else:
        break
score = len(set([x[0] for x in fav]))**2 + sum([x[1] for x in fav])

print(score)

Submission status: Wrong Answer

Solution code:
import collections

n, k = map(int, input().split())
h = []
for i in range(n):
    a, b = map(int, input().split())
    h.append((a,b))
h.sort(key=lambda x: -x[1])
fav=h[0:k]
tb={}
for t,d in sorted(h[k:], key=lambda x: (-x[0], -x[1])):
    tb[t] = tb.get(t, []) + [d]
alt = []
for t,ds in tb.items():
    alt.append((t,max(ds)))
alt.sort(key=lambda x: -x[1])

# print("h: ",h)
# print("fav: ",fav)
# print("alt: ",alt)
while(len(alt) > 0):
    c = collections.Counter([x[0] for x in fav])
    temp = sorted(filter(lambda x: c[x[0]] > 1, fav), key=lambda x: -x[1])
    if(len(temp) == 0):
        break
    t1, d1 = alt.pop(0)
    if(c[t1] > 0):
        continue
    t2, d2 = temp[-1]

    ts = set([x[0] for x in fav])
    if d1 - d2 + (len(ts)+1)**2 - len(ts)**2 > 0:
        fav.remove((t2,d2))
        fav.append((t, d))

score = len(set([x[0] for x in fav]))**2 + sum([x[1] for x in fav])

print(score)

Submission status: Wrong Answer

Solution code:
from collections import defaultdict

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

items = []
for _ in range(n):
    ti, di = map(int, input().split())
    items.append((di, ti))

items.sort(reverse=True)

topics = defaultdict(int)
for _, ti in items[:k]:
    topics[ti] += 1

best_res = sum(v for v, _ in items[:k]) + len(topics) ** 2
curr_res = sum(v for v, _ in items[:k])

i = k - 1
j = k
while i >= 0 and j < n:
    dj, tj = items[j]
    if tj in topics:
        j += 1
        continue
    di, ti = items[i]
    if topics[ti] > 1:
        topics[ti] -= 1
    else:
        del topics[ti]
    topics[tj] += 1
    curr_res -= di
    curr_res += dj
    best_res = max(curr_res + len(topics) ** 2, best_res)
    i -= 1
        
    

print(best_res)

Submission status: Wrong Answer

Solution code:
from itertools import accumulate
N,K = map(int,input().split())
td = [tuple(map(int,input().split())) for i in range(N)]

td = [(x[0]-1,x[1]) for x in td]
td.sort(key=lambda x:-x[1])

t = [-1]*N
for i in range(N):
    t[td[i][0]] = 0
tdK = td[:K]
for i in range(K):
    t[tdK[i][0]] = 1

t_used = [i for i in range(N) if t[i]==1]
t_not_used = [i for i in range(N) if t[i]==0]

d_top = [0]*N
d_top_candi = [0]*N
d_rest = []
d_candi = []

for i in range(K):
    if d_top[td[i][0]] == 0:
        d_top[td[i][0]] = td[i][1]
    else:
        d_rest.append(td[i][1])
for i in range(K,N):
    if t[td[i][0]] == 0:
        if d_top_candi[td[i][0]] == 0:
            d_top_candi[td[i][0]] = 1
            d_candi.append(td[i][1])
l = len(t_used)
s = sum(d_top)
ans = l**2 + s
d_rest.sort()
d_candi.sort(key=lambda x:-x)

d_rest_cum = list(accumulate(d_rest))
d_candi_cum = list(accumulate(d_candi))

if len(d_rest) == 0:
    print(ans)
else:
    ans += d_rest_cum[-1]

    for k in range(min(len(d_rest),len(d_candi))):
        ans = max((l+k+1)**2 + s + d_candi_cum[k] - d_rest_cum[k],ans)

    print(ans)

Submission status: Wrong Answer

Solution code:
from itertools import accumulate
N,K = map(int,input().split())
td = [tuple(map(int,input().split())) for i in range(N)]

td.sort(key=lambda x:-x[1])
td = [(x[0]-1,x[1]) for x in td]
t = [-1]*N
for i in range(N):
    t[td[i][0]] = 0
tdK = td[:K]
for i in range(K):
    t[tdK[i][0]] = 1

t_used = [i for i in range(N) if t[i]==1]
t_not_used = [i for i in range(N) if t[i]==0]

d_top = [0]*N
d_top_candi = [0]*N
d_rest = []
d_candi = []

for i in range(K):
    if t[td[i][0]] == 1:
        if d_top[td[i][0]] == 0:
            d_top[td[i][0]] = td[i][1]
        else:
            d_rest.append(td[i][1])
for i in range(K,N):
    if t[td[i][0]] == 0:
        if d_top_candi[td[i][0]] == 0:
            d_top_candi[td[i][0]] = 1
            d_candi.append(td[i][1])
l = len(t_used)
ans = l**2+sum(d_top)
d_rest.sort()
d_candi.sort(key=lambda x:-x)

d_rest_cum = list(accumulate(d_rest))
d_candi_cum = list(accumulate(d_candi))

if len(d_rest) == 0:
    print(ans)
else:
    ans += d_rest_cum[-1]

    for k in range(min(len(d_rest),len(d_candi))):
        ans = max(ans + (k+1)*(k+1+2*l) + d_candi_cum[k] - d_rest_cum[k],ans)

    print(ans)

Submission status: Wrong Answer

Solution code:
from itertools import accumulate
N,K = map(int,input().split())
td = [tuple(map(int,input().split())) for i in range(N)]

td.sort(key=lambda x:-x[1])
td = [(x[0]-1,x[1]) for x in td]
t = [-1]*N
for i in range(N):
    t[td[i][0]] = 0
for i in range(K):
    t[td[:K][i][0]] = 1

t_used = [i for i in range(N) if t[i]==1]
t_not_used = [i for i in range(N) if t[i]==0]

d_top = [0]*N
d_top_candi = [0]*N
d_rest = []
d_candi = []

for i in range(K):
    if t[td[i][0]] == 1:
        if d_top[td[i][0]] == 0:
            d_top[td[i][0]] = td[i][1]
        else:
            d_rest.append(td[i][1])
for i in range(K,N):
    if t[td[i][0]] == 0:
        if d_top_candi[td[i][0]] == 0:
            d_top_candi[td[i][0]] = 1
            d_candi.append(td[i][1])

ans = len(t_used)**2+sum(d_top)
d_rest.sort()
d_candi.sort(key=lambda x:-x)

d_rest_cum = list(accumulate(d_rest))
d_candi_cum = list(accumulate(d_candi))

if len(d_rest) == 0:
    print(ans)
else:
    ans += d_rest_cum[-1]

    for k in range(min(len(d_rest),len(d_candi))):
        ans = max(ans + (k+1)*(k+1+2*len(t_used)) + d_candi_cum[k] - d_rest_cum[k],ans)

    print(ans)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-

"""
参考：https://img.atcoder.jp/abc116/editorial.pdf
"""

import sys, re
from collections import deque, defaultdict, Counter
from math import sqrt, hypot, factorial, pi, sin, cos, radians
from heapq import heappop, heappush, heapify, heappushpop
from bisect import bisect_left, bisect_right
from itertools import permutations, combinations, product
from operator import itemgetter, mul
from copy import deepcopy
from functools import reduce, partial
from fractions import Fraction
from string import ascii_lowercase, ascii_uppercase, digits

def input(): return sys.stdin.readline().strip()
def ceil(a, b=1): return int(-(-a // b))
def round(x): return int((x*2+1) // 2)
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(): return list(map(int, input().split()))
sys.setrecursionlimit(10 ** 9)
INF = float('inf')
MOD = 10 ** 9 + 7

N, K = MAP()
tdN = [[0]*2 for i in range(N)]
for i in range(N):
    t, d = MAP()
    tdN[i][0] = t
    tdN[i][1] = d

tdN.sort(key=itemgetter(1),reverse=True)
base = 0
bonus = 0
c = Counter()
changes = []
ans = 0
for i in range(K):
    t, d = tdN[i]
    base += d
    if c[t] == 0:
        bonus += 1
    c[t] += 1
    # 交換候補を集める
    if c[t] >= 2:
        heappush(changes, (d, i))
bonus *= bonus
ans = base + bonus

for i in range(K, N):
    t, d = tdN[i]
    # 既に持ってるネタならスキップ
    if c[t] != 0:
        continue
    else:
        # 1つずつ候補から交換してみて、最大値を取る
        if len(changes):
            _, idx = heappop(changes)
            t2, d2 = tdN[idx]
            base -= d2 - d
            bonus = (sqrt(bonus)+1) ** 2
            ans = max(base + bonus, ans)
            c[t2] += 1
            # 今回交換で入れた物が2つ以上ある種類なら交換候補に追加
            if c[t2] >= 2:
                heappush(changes, (d2, idx))
print(int(ans))

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-

"""
参考：https://img.atcoder.jp/abc116/editorial.pdf
"""

import sys, re
from collections import deque, defaultdict, Counter
from math import sqrt, hypot, factorial, pi, sin, cos, radians
from heapq import heappop, heappush, heapify, heappushpop
from bisect import bisect_left, bisect_right
from itertools import permutations, combinations, product
from operator import itemgetter, mul
from copy import deepcopy
from functools import reduce, partial
from fractions import Fraction
from string import ascii_lowercase, ascii_uppercase, digits

def input(): return sys.stdin.readline().strip()
def ceil(a, b=1): return int(-(-a // b))
def round(x): return int((x*2+1) // 2)
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(): return list(map(int, input().split()))
sys.setrecursionlimit(10 ** 9)
INF = float('inf')
MOD = 10 ** 9 + 7

N, K = MAP()
tdN = [[0]*2 for i in range(N)]
for i in range(N):
    t, d = MAP()
    tdN[i][0] = t
    tdN[i][1] = d

tdN.sort(key=itemgetter(1),reverse=True)
base = 0
bonus = 0
c = Counter()
changes = []
ans = 0
for i in range(K):
    t, d = tdN[i]
    base += d
    if c[t] == 0:
        bonus += 1
    c[t] += 1
    if c[t] >= 2:
        heappush(changes, (d, i))
bonus *= bonus
ans = base + bonus

for i in range(K, N):
    t, d = tdN[i]
    # 既に持ってるネタならスキップ
    if c[t] != 0:
        continue
    else:
        # 1つずつ候補から交換してみて、最大値を取る
        if len(changes):
            _, idx = heappop(changes)
            t2, d2 = tdN[idx]
            base -= d2 - d
            bonus = (sqrt(bonus)+1) ** 2
            ans = max(base + bonus, ans)
print(int(ans))

Submission status: Wrong Answer

Solution code:
def getN():
    return int(input())


def getNM():
    return map(int, input().split())


def getList():
    return list(map(int, input().split()))
import heapq

n, k = getNM()
sushi_all = []

for i in range(n):
    sushi_all.append(getList())

sushi_all.sort(reverse=True, key=lambda x: x[1])
#print(sushi_all)

m = len(list(set([x[0] for x in sushi_all[:k]])))
fm = sum([x[1] for x in sushi_all[:k]])
fx = [fm] * k + [0] * (n-k+1)
#print(fx)

heap_used = []
heap_not = []

flag = [0] * (n+1)
for sushi in [x for x in sushi_all[:k]]:
    if flag[sushi[0]] == 1:
        heap_used.append(sushi)
    else:
        flag[sushi[0]]  = 1

heap_used.sort(key=lambda x: -x[1])
#heapq.heapify(sushi_all, key=lambda x: x[1])
#print(heap_used)
for sushi in sushi_all:
    if sum(flag) == k:

        break
    if flag[sushi[0]] == 0:

        m += 1
        fx[m] = fx[m-1] + sushi[1] - heapq.heappop(heap_used)[1]
        flag[sushi[0]] = 1

gx = [fx[i] + i*i if fx[i] != 0 else 0 for i in range(len(fx))]
print(max(gx))


"""
5 2
1 9
1 7
2 6
2 5
3 1
"""
###

Submission status: Wrong Answer

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

td=sorted([tuple(map(int,input().split())) for _ in range(N)],key=lambda x:-x[1])
answer=0
number=[]
afterpop=[]

for i in range(K):
  answer+=td[i][1]
  if td[i][0] not in number:
    number.append(td[i][0])
  else:
    afterpop.append(td[i][1])

answer+=(len(number))**2
    
index=K
tempanswer=answer
while index<N and afterpop:
  if td[i][0] not in number:
    number.append(td[i][0])
    tempanswer=tempanswer-afterpop.pop()-(len(number)-1)**2+(len(number))**2
    answer=max(tempanswer,answer)
  index+=1
                  
print(answer)

Submission status: Wrong Answer

Solution code:
def run(N, K, t, d):
    '''
    おいしさ(d)の昇順でソートする O(NlogN)
    おいしさが大きいもののK個の合計を取得 O(K)
    K個のうちおいしさが少ないものから順にネタの種類が
    増えるように変更していく O(K)
    ネタの種類のうち最大の物の値とインデックスを取得しておく O(K)
    ネタの種類を変えた合計のうち最大のものが答え
    '''
    t_dic = dict(zip(range(N), t))
    d_dic = dict(zip(range(N), d))
    d_dic_sorted = dict(sorted(d_dic.items(), key=lambda x: x[1], reverse=True))
    d_sum = sum(list(d_dic_sorted.values())[:K])
    t_used = {}
    t_cnt = 0
    for i in t:
        t_used[i] = 0
    for i in list(d_dic_sorted.keys())[:K]:
        t_used[t_dic[i]] += 1
        if t_used[t_dic[i]] == 1:
            t_cnt += 1
    t_sum = t_cnt ** 2
    max_sum = d_sum + t_sum
    del_list = list(d_dic_sorted.keys())[:K]
    for i in list(d_dic_sorted.keys())[K:]:
        di = d_dic_sorted[i]
        ti = t_dic[i]
        # ネタの種類が増えるように追加
        if t_used[ti] == 0:
            t_used[ti] += 1
            t_cnt += 1
            t_sum += 2*t_cnt - 1
            d_sum += di
            # ネタが重複しているものを削除
            flg = True
            while len(del_list) > 0 and flg:
                del_i = del_list.pop()
                del_di = d_dic_sorted[del_i]
                del_ti = t_dic[del_i]
                if t_used[del_ti] > 1:
                    t_used[del_ti] -= 1
                    d_sum -= del_di
                    max_sum = max(max_sum, d_sum + t_sum)
                    flg = False
    return max_sum


def main():
    N, K = map(int, input().split())
    t = []
    d = []
    for _ in range(N):
        ti, di = map(int, input().split())
        t.append(ti)
        d.append(di)
    print(run(N, K, t, d))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
def run(N, K, t, d):
    '''
    おいしさ(d)の昇順でソートする O(NlogN)
    おいしさが大きいもののK個の合計を取得 O(K)
    K個のうちおいしさが少ないものから順にネタの種類が
    増えるように変更していく O(K)
    ネタの種類のうち最大の物の値とインデックスを取得しておく O(K)
    ネタの種類を変えた合計のうち最大のものが答え
    '''
    t_dic = dict(zip(range(N), t))
    d_dic = dict(zip(range(N), d))
    d_dic_sorted = dict(sorted(d_dic.items(), key=lambda x: x[1], reverse=True))
    d_sum = sum(list(d_dic_sorted.values())[:K])
    t_used = {}
    t_cnt = 0
    for i in t:
        t_used[i] = 0
    for i in list(d_dic_sorted.keys())[:K]:
        t_used[t_dic[i]] += 1
        if t_used[t_dic[i]] == 1:
            t_cnt += 1
    t_sum = t_cnt ** 2
    max_sum = d_sum + t_sum
    del_list = list(d_dic_sorted.keys())[:K]
    for i in list(d_dic_sorted.keys())[K:]:
        di = d_dic_sorted[i]
        ti = t_dic[i]
        # ネタの種類が増えるように追加
        if t_used[ti] == 0:
            t_used[ti] += 1
            t_cnt += 1
            t_sum += 2*t_cnt - 1
            d_sum += di
            # ネタが重複しているものを削除
            while len(del_list) > 0:
                del_i = del_list.pop()
                del_di = d_dic_sorted[del_i]
                del_ti = t_dic[del_i]
                if t_used[del_ti] > 1:
                    t_used[del_ti] -= 1
                    d_sum -= del_di
                    max_sum = max(max_sum, d_sum + t_sum)
                    break
    return max_sum


def main():
    N, K = map(int, input().split())
    t = []
    d = []
    for _ in range(N):
        ti, di = map(int, input().split())
        t.append(ti)
        d.append(di)
    print(run(N, K, t, d))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
def run(N, K, t, d):
    '''
    おいしさ(d)の昇順でソートする O(NlogN)
    おいしさが大きいもののK個の合計を取得 O(K)
    K個のうちおいしさが少ないものから順にネタの種類が
    増えるように変更していく O(K)
    ネタの種類のうち最大の物の値とインデックスを取得しておく O(K)
    ネタの種類を変えた合計のうち最大のものが答え
    '''
    t_dic = dict(zip(range(N), t))
    d_dic = dict(zip(range(N), d))
    d_dic_sorted = dict(sorted(d_dic.items(), key=lambda x: x[1], reverse=True))
    d_sum = sum(list(d_dic_sorted.values())[:K])
    t_used = {}
    t_cnt = 0
    for i in t:
        t_used[i] = 0
    for i in list(d_dic_sorted.keys())[:K]:
        t_used[t_dic[i]] += 1
        if t_used[t_dic[i]] == 1:
            t_cnt += 1
    t_sum = t_cnt ** 2
    max_sum = d_sum + t_sum
    del_list = list(d_dic_sorted.keys())[:K]
    for i in list(d_dic_sorted.keys())[K:]:
        di = d_dic_sorted[i]
        ti = t_dic[i]
        # ネタの種類が増えるように追加
        if t_used[ti] == 0:
            t_used[ti] += 1
            t_cnt += 1
            t_sum += 2*t_cnt - 1
            d_sum += di
            # ネタが重複しているものを削除
            while len(del_list) > 0:
                del_i = del_list.pop()
                del_di = d_dic_sorted[del_i]
                del_ti = t_dic[del_i]
                if t_used[del_ti] > 1:
                    t_used[del_ti] -= 1
                    d_sum -= del_di
                    del_i -= 1
                    max_sum = max(max_sum, d_sum + t_sum)
                    break
                else:
                    del_i -= 1
                    continue
    return max_sum


def main():
    N, K = map(int, input().split())
    t = []
    d = []
    for _ in range(N):
        ti, di = map(int, input().split())
        t.append(ti)
        d.append(di)
    print(run(N, K, t, d))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
def solver(array):
    sum_d = 0
    t_array = []
    for i in range(len(array)):
        sum_d += array[i][1]
        t_array.append(array[i][0])

    return sum_d + len(set(t_array)) * len(set(t_array))


n, k = map(int, input().split())
td_array = [[int(x) for x in input().split()] for y in range(n)]
various_array = set([x for x, y in td_array])
satisfy_array = [y for x, y in td_array]

# print(various_array, satisfy_array)
# print(td_array)

td_dict = {x:[] for x in set([y for y, z in td_array])}
for t, d in td_array:
    td_dict[t].append(d)
td_dict = {key: sorted(value) for key, value in td_dict.items()}
key_array = [x for x in td_dict.keys()]
prepared_array = []
while True:
    tmp_array = []
    for key in key_array:
        if td_dict[key]:
            tmp_array.append([key, td_dict[key].pop(-1)])
        else:
            pass
    prepared_array += sorted(tmp_array, key=lambda x:x[1], reverse=True)
    for value in td_dict.values():
        if value:
            break
    else:
        break

# print(prepared_array)
# td_array = sorted(td_array, key=lambda x: x[1], reverse=True)
res_array = prepared_array[:k]
for td in prepared_array[k:]:
    if solver(res_array) < solver(res_array[:-1] + [td]):
        res_array = res_array[:-1] + [td]
print(solver(res_array))

Submission status: Wrong Answer

Solution code:
init = input().split()
N = int(init[0])
K = int(init[1])

values = []

category = set()
for _ in range(N):
    kind, value = map(int, input().split())
    values.append([kind, value])


values = sorted(values, key=lambda x:x[1], reverse=True)

result = 0
notfirsts = []
for v, p in values[:K]:
    if v not in category:
        category.add(v)
    else:
        notfirsts.append(p)
    result += p

base_result = len(category)**2 + result + 1 
max_result = base_result
print(max_result, category, notfirsts)


for v, p in values[K:]:
    if not notfirsts:
        break
    if v not in category:
        base_result += (2 * len(category)) + 1 + p - notfirsts.pop() 
        category.add(v)
        max_result = max(max_result, base_result)
        print(max_result)

Submission status: Wrong Answer

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

explored = []
X = []
Y = []
for i in range(N) :
    if L[i][0] not in explored :
        explored.append(L[i][0])
        Y.append(L[i][1])
    else :
        X.append(L[i][1])

ans = 0
for i in range(len(Y)) :
    tmp = sum(Y[:i+1]) + sum(X[:K-(i+1)]) + len(Y[:i+1])**2
    #print(sum(Y[:i+1]), sum(X[:K-(i+1)]), len(Y[:i+1])**2)
    if ans < tmp :
        ans = tmp
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq as hq
N, K = map(int, input().split())
Sushi = [None for _ in range(N)]
TotalType = set()
for _ in range(N):
    t, d = map(int, input().split())
    Sushi[_] = [d, t]
Sushi.sort(reverse = True)

ChosenType = set()
Double = []
hq.heapify(Double)
NotChosen = Sushi[K:]
BasePoint = 0
for i in range(K):
    BasePoint += Sushi[i][0]
    if Sushi[i][1] in ChosenType:
        hq.heappush(Double, Sushi[i])
    else: ChosenType |= {Sushi[i][1]}

MaxHappiness = BasePoint + len(ChosenType)**2
while Double and NotChosen:
    TrushPoint, TrushType = hq.heappop(Double)
    BasePoint -= TrushPoint
    for add in NotChosen:
        if add[1] in ChosenType:
            NotChosen.remove(add)
        else:
            BasePoint += add[0]
            ChosenType |= {add[1]}
            break
    MaxHappiness = max(MaxHappiness, BasePoint + len(ChosenType)**2)
print(MaxHappiness)

Submission status: Wrong Answer

Solution code:
from operator import itemgetter
import sys
input = sys.stdin.readline

def inpl(): return [int(i) for i in input().split()]

N,K = inpl()
ans = 0
n0 = []
n1 = [0]*(N+1)

for i in range(N):
    ti,di = inpl()
    if(n1[ti]<di):
        if(n1[ti]!=0):
            n0.append(n1[ti])
        n1[ti]=di
    else:
        n0.append(di)

n0.sort(reverse = True)
n1.sort(reverse = True)

for i in range(min(K+1,len(n1)+1)):
    if(K-i <= len(n0) and 0<=K-i):
        total = sum(n1[:i]+n0[:K-i])+i*i
        if(ans<total):
            ans = total

print(ans)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
 
n, k = map(int, input().split())
 
sushi_list = []
neta_set = set()
for i in range(n):
    t, d = map(int, input().split())
    sushi_list.append((t, d))
    neta_set.add(t)
 
sushi_list.sort(key=lambda tup: tup[1])
 
max_sushi_list, rest_of_sushi = sushi_list[-k:], sushi_list[:-k]
max_neta_set = {sushi[0] for sushi in max_sushi_list}
 
iMax = len(neta_set)
iMin = len(max_neta_set)
 
max_score = iMin**2 + sum([sushi[1] for sushi in max_sushi_list])
 
 
for nNeta in range(iMin+1, iMax + 1, 1):
 
    # Find sushi having minimum oishisa and being not unique for neta
    r = len(max_sushi_list)
    swap_index = -1
    for i in range(r):
        if max_sushi_list[i][0] in [neta for neta,oishisa in max_sushi_list[:i]+max_sushi_list[i+1:]]:
            swap_index = i
            break
 
    # Find sushi having maximum oishisa and unique for neta
    r = len(rest_of_sushi)
    for i in range(r):
        if not rest_of_sushi[r-i-1][0] in max_neta_set:
            rest_of_sushi.append(max_sushi_list.pop(swap_index))
            max_sushi_list.append(rest_of_sushi.pop(r-i-1))
            break
 
    score = nNeta**2 + sum([sushi[1] for sushi in max_sushi_list])
    if score > max_score:
      max_score = score
    sushi_list.sort(key=lambda tup: tup[1])
    max_neta_set = {sushi[0] for sushi in max_sushi_list}
 
print(max_score)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop
#import collections
n, k = map(int, input().split())
c = []
for i in range(n):#h:高さ
    c.append([int(m) for m in input().split()])
c.sort(key=lambda x:(x[0], x[1]))
f = sorted(c, key=lambda x:(x[1], x[0]), reverse=True)
donyoku = 0
donset = set()
#pq = PriorityQueue()
notfirst =[]
for i in range(k):
    if f[i][0] not in donset:
        donset.add(f[i][0])
    else:
        heappush(notfirst, f[i][1])
    donyoku += f[i][1]
donyoku += len(donset)**2

candidate = []
candidateset = set()
for i in range(k, n):
    if f[i][0] not in donset:
        candidate.append(-f[i][1])#最小のものから取り出されるのでマイナスをつけることによって最大のものから取り出せるようにしている
        candidateset.add(f[i][0])

satili = [donyoku]
donyoku -= len(donset)**2
for i in range(len(candidateset)):
    if notfirst == []:
        break
    else:
        donyoku -= heappop(notfirst)
        donyoku += -heappop(candidate)#最小のものから取り出されるのでマイナスをつけることによって最大のものから取り出せるようにしている
        donyoku += (len(donset) + i + 1) ** 2
        satili.append(donyoku)
    donyoku -= (len(donset) + i + 1) ** 2




print(max(satili))