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

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

    if not s[0] in waiting_sushi:
        waiting_sushi[s[0]] = []
    l = waiting_sushi[s[0]]
    heapq.heappush(l, (s[1], s[0]))

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

    if not target:
        break

    # 入れ替えして比較
    s = heapq.heappop(waiting_sushi[s_kind])
    val = max_pt - target[1] + s[1] + (kind_cnt + 1) ** 2
    if val > max_pt:
        max_pt = val
        kind_counter[target[0]] -= 1
        kind_cnt += 1

    idx += 1
print(max_pt)

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

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

    if not s[0] in waiting_sushi:
        waiting_sushi[s[0]] = []
    l = waiting_sushi[s[0]]
    heapq.heappush(l, (s[1], s[0]))

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

    if not target:
        break

    # 入れ替えして比較
    s = heapq.heappop(waiting_sushi[s_kind])
    val = max_pt - target[1] + s[1] + (kind_cnt + 1) ** 2
    if val > max_pt:
        max_pt = val
        kind_counter[target[0]] -= 1

    idx += 1
print(max_pt)

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 = []
for i in range(n):
  t,d = td[i]
  if t not in eated_kind:
    eated_kind.append(t)
    eated_kind_num += 1
    eated.append([t,d])
    
    ###tがk種類存在したら、残りのtを扱うことはない
    ###またdも最大化されているため、これ以降のtdを使うことはない
    if eated_kind_num == k:
      td = td[:i+1]
      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:
import copy
from collections import defaultdict
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: x[1], reverse=True)

cur = TD[:K]
base = sum([x[1] for x in cur])
variety = defaultdict(int)
for t, d in cur:
    variety[t] += 1
bonus = len(variety) ** 2

ans = base + bonus
for i in range(K, N):
    new_cur = copy.copy(cur)
    del new_cur[-1]
    new_cur.append(TD[i])
    new_var = copy.copy(variety)
    new_var[cur[-1][0]] -= 1
    if new_var[cur[-1][0]] == 0:
        del new_var[cur[-1][0]]
    new_var[TD[i][0]] += 1
    new_ans = base - cur[-1][1] + TD[i][1] - len(variety) ** 2 + len(new_var) ** 2
    if new_ans > ans:
        ans = new_ans
        cur = new_cur
        variety = new_var
    print(ans, cur, variety)

print(ans)

Submission status: Wrong Answer

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

# 集合の中の種類数を出す


def v(l):
    t = set([])
    for i in range(len(l)):
        t.add(l[i][2])
    return len(t)


def calcNext(base, netas, types, index, sumNext, backIndex):
    removeTarget = 0
    for i in range(backIndex, len(base) + 1):
        t = base[-i][2]
        if types[t] > 1:
            removeTarget = len(base)-i
            types[t] -= 1
            sumNext -= base[-i][1]
            backIndex = i + 1
            break
    #next = []
    # for i in range(len(base)):
    #     if i == removeTarget:
    #         continue
    #     next.append(base[i])
    nextIndex = index
    found = False
    for i in range(index, len(netas)):
        nextIndex = i
        if types[netas[i][2]] == 0:
            base[removeTarget] = netas[i]
            types[netas[i][2]] = 1
            sumNext += netas[i][1]
            found = True
            break
    return (base, nextIndex, sumNext, backIndex, found)


def score(l):
    sum = 0
    for i in l:
        sum += i[1]
    va = v(l)
    sum += va*va
    return sum


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    netas = [(-d[i], d[i], t[i]) for i in range(N)]
    netas.sort()
    S = []
    T = collections.defaultdict(int)
    sum1 = 0
    for i in range(K):
        S.append(netas[i])
        sum1 += netas[i][1]
        T[netas[i][2]] += 1
    a = v(S)
    sMax = sum1 + a*a
    lastSum = sum1
    index = K
    backIndex = 1
    for t in range(a+1, N):
        S, index, lastSum, backIndex, found = calcNext(
            S, netas, T, index, lastSum, backIndex)
        if not found:
            break
        sMax = max(lastSum + t*t, sMax)
    print(sMax)
    return


# Generated by 1.1.3 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:
#!/usr/bin/env python3
import sys
import heapq
import collections

# 集合の中の種類数を出す


def v(l):
    t = set([])
    for i in range(len(l)):
        t.add(l[i][2])
    return len(t)


def calcNext(base, netas, types, index, sumNext, backIndex):
    removeTarget = 0
    for i in range(backIndex, len(base) + 1):
        t = base[-i][2]
        if types[t] > 1:
            removeTarget = len(base)-i
            types[t] -= 1
            sumNext -= base[-i][1]
            backIndex = i + 1
            break
    #next = []
    # for i in range(len(base)):
    #     if i == removeTarget:
    #         continue
    #     next.append(base[i])
    nextIndex = index
    for i in range(index, len(netas)):
        nextIndex = i
        if types[netas[i][2]] == 0:
            base[removeTarget] = netas[i]
            types[netas[i][2]] = 1
            sumNext += netas[i][1]
            break
    return (base, nextIndex, sumNext, backIndex)


def score(l):
    sum = 0
    for i in l:
        sum += i[1]
    va = v(l)
    sum += va*va
    return sum


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    netas = [(-d[i], d[i], t[i]) for i in range(N)]
    netas.sort()
    maxT = v(netas)
    S = []
    C = [(0, []) for i in range(maxT+1)]
    T = collections.defaultdict(int)
    sum1 = 0
    for i in range(K):
        S.append(netas[i])
        sum1 += netas[i][1]
        T[netas[i][2]] += 1
    a = v(S)
    C[a] = (sum1, S)
    index = K
    backIndex = 1
    for t in range(a+1, maxT+1):
        nl, index, nsum, backIndex = calcNext(
            C[t-1][1], netas, T, index, C[t-1][0], backIndex)
        C[t] = (nsum, nl)
    sMax = 0
    for i in range(a, maxT+1):
        sMax = max(sMax, C[i][0] + i*i)
    print(sMax)
    return


# Generated by 1.1.3 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:
#!/usr/bin/env python3
import sys
import heapq
import collections

# 集合の中の種類数を出す


def v(l):
    t = set([])
    for i in range(len(l)):
        t.add(l[i][2])
    return len(t)


def calcNext(base, netas, types, index, sumNext, backIndex):
    removeTarget = 0
    for i in range(backIndex, len(base) + 1):
        t = base[-i][2]
        if types[t] > 1:
            removeTarget = len(base)-i
            types[t] -= 1
            sumNext -= base[-i][1]
            backIndex = i + 1
            break
    next = []
    for i in range(len(base)):
        if i == removeTarget:
            continue
        next.append(base[i])
    nextIndex = index
    for i in range(index, len(netas)):
        nextIndex = i
        if types[netas[i][2]] == 0:
            next.append(netas[i])
            types[netas[i][2]] = 1
            sumNext += netas[i][1]
            break
    return (next, nextIndex, sumNext, backIndex)


def score(l):
    sum = 0
    for i in l:
        sum += i[1]
    va = v(l)
    sum += va*va
    return sum


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    netas = [(-d[i], d[i], t[i]) for i in range(N)]
    netas.sort()
    maxT = v(netas)
    S = []
    C = [(0, []) for i in range(maxT+1)]
    T = collections.defaultdict(int)
    sum1 = 0
    for i in range(K):
        S.append(netas[i])
        sum1 += netas[i][1]
        T[netas[i][2]] += 1
    a = v(S)
    C[a] = (sum1, S)
    index = K
    backIndex = 1
    for t in range(a+1, maxT+1):
        nl, index, nsum, backIndex = calcNext(
            C[t-1][1], netas, T, index, C[t-1][0], backIndex)
        C[t] = (nsum, nl)
    sMax = 0
    for i in range(a, maxT+1):
        sMax = max(sMax, C[i][0] + i*i)
    print(sMax)
    return


# Generated by 1.1.3 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:
#!/usr/bin/env python3
import sys
import heapq
import collections

# 集合の中の種類数を出す


def v(l):
    t = set([])
    for i in range(len(l)):
        t.add(l[i][2])
    return len(t)


def calcNext(base, netas, types, index, sumNext):
    removeTarget = 0
    for i in range(1, len(base) + 1):
        t = base[-i][2]
        if types[t] > 1:
            removeTarget = len(base)-i
            types[t] -= 1
            sumNext -= base[-i][1]
            break
    next = []
    for i in range(len(base)):
        if i == removeTarget:
            continue
        next.append(base[i])
    nextIndex = index
    for i in range(index, len(netas)):
        nextIndex = i
        if types[netas[i][2]] == 0:
            next.append(netas[i])
            types[netas[i][2]] = 1
            sumNext += netas[i][1]
            break
    return (next, nextIndex, sumNext)


def score(l):
    sum = 0
    for i in l:
        sum += i[1]
    va = v(l)
    sum += va*va
    return sum


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    netas = [(-d[i], d[i], t[i]) for i in range(N)]
    netas.sort()
    maxT = v(netas)
    S = []
    C = [(0, []) for i in range(maxT+1)]
    T = collections.defaultdict(int)
    sum1 = 0
    for i in range(K):
        S.append(netas[i])
        sum1 += netas[i][1]
        T[netas[i][2]] += 1
    a = v(S)
    C[a] = (sum1, S)
    index = K
    for t in range(a+1, maxT+1):
        nl, index, nsum = calcNext(C[t-1][1], netas, T, index, C[t-1][0])
        C[t] = (nsum, nl)
    sMax = 0
    for i in range(a, maxT+1):
        sMax = max(sMax, C[i][0] + i*i)
    print(sMax)
    return


# Generated by 1.1.3 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:
#!/usr/bin/env python3
import sys
import heapq


def solve(N: int, K: int, t: "List[int]", d: "List[int]"):
    heaps = {}
    for i in range(N):
        if t[i] not in heaps:
            heaps[t[i]] = []
            heapq.heapify(heaps[t[i]])
        heap = heaps[t[i]]
        heapq.heappush(heap, (-d[i], d[i]))
    used = set([])
    current = 0
    for i in range(K):
        maxIndex = -1
        max = -1
        for d in heaps:
            heap = heaps[d]
            if len(heap) == 0:
                continue
            next = heap[0][1]
            if d not in used:
                next += (len(used)+1) * (len(used)+1) - len(used)*len(used)
            if max < next:
                max = next
                maxIndex = d
        if maxIndex != -1:
            used.add(maxIndex)
            current += max
            heapq.heappop(heaps[maxIndex])
            # print(f"use {maxIndex}")
    print(current)
    return


# Generated by 1.1.3 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:
N,K = map(int,input().split())
TD = [list(map(int,input().split())) for _ in range(N)]
T,D = [],[]
for t,d in TD:
    T.append(t)
    D.append(d)
import collections
tt = collections.Counter(T)

bonus=len(tt)*len(tt)

TD.sort()
a = bonus+sum(D[-K:])

i = 1
ans = 0
res = 0
total = 0

for t,d in TD:
    if t > i:
        i = t
        total += res
        res = 0
    if t == i:
        if res < d:
            res = d
            if TD[-1] == [t,d]:
                total += res
res = total+bonus
TD.sort(reverse=True,key=lambda x: x[1])
tl = []
total = 0
for t,d in TD[:K]:
    tl.append(t)
    total += d
bonus = len(collections.Counter(tl))**2
total += bonus
print(max(res,total))

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
TD = [list(map(int,input().split())) for _ in range(N)]
T,D = [],[]
for t,d in TD:
    T.append(t)
    D.append(d)
import collections
tt = collections.Counter(T)

bonus=len(tt)*len(tt)
TD.sort(reverse=True,key=lambda x: x[1])
a = bonus+sum(D[-K:])

i = 1
ans = 0
res = 0
total = 0

for t,d in TD:
    if t > i:
        i = t
        total += res
        res = 0
    if t == i:
        if res < d:
            res = d
            if TD[-1] == [t,d]:
                total += res
res = total+bonus

tl = []
total = 0
for t,d in TD[:K]:
    tl.append(t)
    total += d
bonus = len(collections.Counter(tl))**2
total += bonus
print(max(res,total))

Submission status: Wrong Answer

Solution code:
import heapq

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

for _ in range(N):
    (t,d)=map(int,input().split())
    heapq.heappush(L,(-d,t))

LL=[]
sum=0
used=set()

for i in range(K):
    d,t=heapq.heappop(L)
    if t in used:
        heapq.heappush(LL,(d,t))
    used.add(t)
    sum+=-d

ans=len(used)**2+sum

while L and LL:
    d,t=heapq.heappop(L)
    if t not in used:
        used.add(t)
        dd,_=heapq.heappop(LL)
        sum=sum+dd+(-d)
        ans=max(ans,len(used)**2+sum)

print(ans)

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
x = []
for i in range(n):
    x.append(list(map(int, input().split())))
x = sorted(x, key=lambda x: x[1])[::-1]
y = {}
z = [0]
for i in range(k):
    z[-1] += x[i][1]
    if x[i][0] in y.keys():
        y[x[i][0]] += 1
    else:
        y[x[i][0]] = 1
z[-1] += len(y.keys()) ** 2

j = 0
while(True):
    z.append(z[-1])
    flag0 = False
    for i in range(k - 1 - j, -1, -1):
        if y[x[i][0]] > 1:
            y[x[i][0]] -= 1
            z[-1] -= x[i][1]
            flag0 = True
            break
    flag1 = False
    for i in range(k + j, n):
        if x[i][0] not in y.keys():
            y[x[i][0]] = 1
            z[-1] += x[i][1]
            flag1 = True
            break
    if flag0 and flag1:
        z[-1] += 2 * len(y.keys()) - 1
    else:
        del z[-1]
        break
    if k + j + 1 == n or len(y.keys()) == k or z[-1] == z[-2]:
        break
    else:
        j += 1
print(max(z))

Submission status: Wrong Answer

Solution code:
n, k = map(int, input().split())
x = []
for i in range(n):
    x.append(list(map(int, input().split())))
x = sorted(x, key=lambda x: x[1])[::-1]
y = {}
z = [0]
for i in range(k):
    z[-1] += x[i][1]
    if x[i][0] in y.keys():
        y[x[i][0]] += 1
    else:
        y[x[i][0]] = 1
z[-1] += len(y.keys()) ** 2

j = 0
while(True):
    z.append(z[-1])
    for i in range(k - 1 - j, -1, -1):
        if y[x[i][0]] > 1:
            y[x[i][0]] -= 1
            z[-1] -= x[i][1]
            break
    for i in range(k + j, n):
        if x[i][0] not in y.keys():
            y[x[i][0]] = 1
            z[-1] += x[i][1]
            break
    z[-1] += 2 * len(y.keys()) - 1
    if len(y.keys()) == k or z[-1] == z[-2]:
        break
    else:
        j += 1
print(max(z))

Submission status: Wrong Answer

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

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

X = []
Y = []  # 一番美味しい寿司
kind = set()
for i in range(N):
    t, d = TD[i]
    if t in kind:
        X.append(d)
    else:
        Y.append(d)
        kind.add(t)

L = len(Y)
ans, rans = 0, 0
for i in range(1, L + 1):  # Yからi個選択 Xからは(K-i)個
    if len(X) < K - i:
        continue
    ans = sum(Y[0:i]) + (i**2) + sum(X[0:K - i])
    rans = max(rans, ans)
print(rans)

Submission status: Wrong Answer

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

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

sushi0 = []
sushi1 = []

for i in range(N):
    if i == 0:
        sushi1.append(TD[0][1])
    else:
        if TD[i][0] == TD[i - 1][0]:
            sushi0.append(TD[i][1])
        else:
            sushi1.append(TD[i][1])

sushi0 = sorted(sushi0, reverse=True)
sushi1 = sorted(sushi1, reverse=True)

L = len(sushi1)
ans, rans = 0, 0
for i in range(1, L + 1):  # sushi1からi個選択 sushi0からは(K-i)個
    if len(sushi0) < K - i:
        continue
    ans = sum(sushi1[0:i]) + (i**2) + sum(sushi0[0:K - i])
    rans = max(rans, ans)
print(rans)

Submission status: Wrong Answer

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

TD = sorted(TD, key=lambda x: x[0])

sushi0 = []
sushi1 = []

for i in range(N):
    if i == 0:
        sushi1.append(TD[0][1])
    else:
        if TD[i][0] == TD[i - 1][0]:
            sushi0.append(TD[i][1])
        else:
            sushi1.append(TD[i][1])

sushi0 = sorted(sushi0, reverse=True)
sushi1 = sorted(sushi1, reverse=True)

L = len(sushi1)
ans, rans = 0, 0
for i in range(1, L + 1):  # sushi1からi個選択 sushi0からは(K-i)個
    if len(sushi0) < K - i:
        continue
    ans = sum(sushi1[0:i]) + (i**2) + sum(sushi0[0:K - i])
    rans = max(rans, ans)
print(rans)

Submission status: Wrong Answer

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

TD = sorted(TD, key=lambda x: x[0])

sushi0 = []
sushi1 = []

for i in range(N):
    if i == 0:
        sushi1.append(TD[0][1])
    else:
        if TD[i][0] == TD[i - 1][0]:
            sushi0.append(TD[i][1])
        else:
            sushi1.append(TD[i][1])

sushi0 = sorted(sushi0, reverse=True)
sushi1 = sorted(sushi1, reverse=True)

L = len(sushi1)
ans, rans = 0, 0
for i in range(1, L + 1):  # sushi1からi個選択 sushi0からは(K-i)個
    ans = sum(sushi1[0:i]) + (i**2) + sum(sushi0[0:K - i])
    rans = max(rans, ans)
print(rans)

Submission status: Wrong Answer

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

TD = sorted(TD, key=lambda x: x[0])

sushi0 = []
sushi1 = []

kind = 1
for i in range(N):
    if i == 0:
        sushi1.append(TD[0][1])
    else:
        if TD[i][0] == TD[i - 1][0]:
            sushi0.append(TD[i][1])
        else:
            kind += 1
            sushi1.append(TD[i][1])

sushi0 = sorted(sushi0, reverse=True)
sushi1 = sorted(sushi1, reverse=True)

L = len(sushi1)
ans = 0
realans = 0
for i in range(L + 1):  # sushi1からi個選択 sushi0からは(K-i)個
    ans = sum(sushi1[0:i]) + (i**2) + sum(sushi0[0:K - i])
    realans = max(realans, ans)
print(realans)

Submission status: Wrong Answer

Solution code:
from heapq import*
n,k=map(int,input().split())
l=[tuple(map(int,input().split()))for i in range(n)]
heapify(sorted(l,key=lambda x:x[1])[::-1])
print(l[:k])
print(l[k:])

Submission status: Wrong Answer

Solution code:
from collections import Counter
N, K =  map(int, (input().split(' ')))
sushi = [tuple(map(int, input().split(' '))) for _ in range(N)]
all_sushi = sorted(sushi, key=lambda x: -x[1])
default_sushi = all_sushi[:K]
amari = all_sushi[K:]
counter = Counter([i[0] for i in default_sushi])
result = []
def calc_value(sushi):
    return sum([i[1] for i in sushi]) + len(counter)**2

def change_sushi(sushi):
    for s in sushi[::-1]:
        if counter[s[0]] >= 2:
            for i, a in enumerate(amari):
                if a[0] not in counter:
                    counter[s[0]] -= 1
                    counter[a[0]] += 1
                    #amari = amari[:i] + amari[i+1:]
                    sushi = [a] + sushi[:-1]
                    return sushi
    return []
            
    
def get_sushi(sushi):
    result.append(calc_value(sushi))
    changed_sushi = change_sushi(sushi)
    if not changed_sushi:
        print(max(result))
        return
    get_sushi(changed_sushi)

get_sushi(default_sushi)

Submission status: Wrong Answer

Solution code:
from collections import Counter
N, K =  map(int, (input().split(' ')))
sushi = [tuple(map(int, input().split(' '))) for _ in range(N)]
all_sushi = sorted(sushi, key=lambda x: -x[1])
default_sushi = all_sushi[:K]
amari = all_sushi[K:]
counter = Counter([i[0] for i in default_sushi])
result = []
def calc_value(sushi):
    return sum([i[1] for i in sushi]) + len(set([i[0] for i in sushi]))**2

def change_sushi(sushi):
    for s in sushi[::-1]:
        if counter[s[0]] >= 2:
            for i, a in enumerate(amari):
                if a[0] not in counter:
                    counter.subtract(s)
                    counter.update([a[0]])
                    new_sushi = a
                    del amari[i]
                    #amari = amari[:i] + amari[i+1:]
                    sushi = [new_sushi] + sushi[:-1]
                    return sushi
    return []
            
    
def get_sushi(sushi):
    result.append(calc_value(sushi))
    changed_sushi = change_sushi(sushi)
    if not changed_sushi:
        print(max(result))
        return
    get_sushi(changed_sushi)

get_sushi(default_sushi)

Submission status: Wrong Answer

Solution code:
from collections import Counter
N, K =  map(int, (input().split(' ')))
sushi = [tuple(map(int, input().split(' '))) for _ in range(N)]
all_sushi = sorted(sushi, key=lambda x: -x[1])
print(sushi)
default_sushi = all_sushi[:K]
amari = all_sushi[K:]
counter = Counter([i[0] for i in default_sushi])
result = []
def calc_value(sushi):
    return sum([i[1] for i in sushi]) + len(set([i[0] for i in sushi]))**2

def change_sushi(sushi):
    for s in sushi[::-1]:
        if counter[s[0]] >= 2:
            for i, a in enumerate(amari):
                if a[0] not in counter:
                    counter.subtract(s)
                    counter.update([a[0]])
                    new_sushi = a
                    del amari[i]
                    #amari = amari[:i] + amari[i+1:]
                    sushi = [new_sushi] + sushi[:-1]
                    return sushi
    return []
            
    
def get_sushi(sushi):
    result.append(calc_value(sushi))
    changed_sushi = change_sushi(sushi)
    if not changed_sushi:
        print(max(result))
        return
    get_sushi(changed_sushi)

get_sushi(default_sushi)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
best = [0] * N
not_best = [0] * N
for _ in range(N):
    t, d = map(int, input().split())
    if best[t-1] != 0:
        not_best.append(min(d, best[t-1]))
    best[t-1] = max(d, best[t-1])
best.sort(reverse=True)
not_best.sort(reverse=True)
acc_best = [0] * (N + 1)
acc_not_best = [0] * (N + 1)
for i in range(N):
    acc_best[i+1] = acc_best[i] + best[i]
for i in range(N):
    acc_not_best[i+1] = acc_not_best[i] + not_best[i]
cand = []
for x in range(1, K+1):
    cand.append(acc_best[x] + acc_not_best[K-x] + x**2)
print(max(cand))

Submission status: Wrong Answer

Solution code:
import collections

def maxsushi(sushi_list,eaten_sushi):
    max_val = [0,0]
    for i in range(len(sushi_list)):
        if(sushi_list[i][1] not in eaten_sushi):
            tmp = sushi_list[i]
            max_val = sushi_list[i]
            break
    return max_val

N,K = map(int,input().split())
sushi = []
umasa = []
for _ in range(N):
    i,j = map(int,input().split())
    sushi.append(i)
    umasa.append(j)
l = list(zip(umasa,sushi))
l = sorted(l,reverse=True)
sushi_list = list(l)
eat_sushi = sushi_list[0:K]
eaten_sushi = set()
double_sushi = []
for i in range(K):
    if(sushi_list[i][1] not in eaten_sushi):
        eaten_sushi.add(sushi_list[i][1])
    else:
        double_sushi.append((sushi_list[i]))
manzoku = 0
sushi_list = sushi_list[K:N]
for i in range(len(eat_sushi)):
    manzoku += eat_sushi[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku
flg = 0
while double_sushi:
    if(flg == 1):
        break
    else:
        tmp1 = double_sushi.pop()
        max_val = [0,0]
        for i in range(len(sushi_list)):
            if(sushi_list[i][1] not in eaten_sushi):
                tmp = sushi_list[i]
                max_val = sushi_list[i]
                break
            flg = 1
    tmp2 = max_val
    if (flg == 1):
        tmp2 =[0,0]
    if(tmp2[0] == 0 and tmp2[1] == 0):
        break
    eaten_sushi.add(tmp2[1])
    tmp_manzoku = tmp_manzoku + tmp2[0] - tmp1[0] + len(eaten_sushi) ** 2 - (len(eaten_sushi) - 1) ** 2
    if (tmp_manzoku > manzoku):
        manzoku = tmp_manzoku
print(manzoku)

Submission status: Wrong Answer

Solution code:
import collections

def maxsushi(sushi_list,eaten_sushi):
    max_val = [0,0]
    for i in range(len(sushi_list)):
        if(sushi_list[i][1] not in eaten_sushi):
            tmp = sushi_list[i]
            max_val = sushi_list[i]
            break
    return max_val

N,K = map(int,input().split())
sushi = []
umasa = []
for _ in range(N):
    i,j = map(int,input().split())
    sushi.append(i)
    umasa.append(j)
l = list(zip(umasa,sushi))
l = sorted(l,reverse=True)
sushi_list = list(l)
eat_sushi = sushi_list[0:K]
eaten_sushi = set()
double_sushi = []
for i in range(K):
    if(sushi_list[i][1] not in eaten_sushi):
        eaten_sushi.add(sushi_list[i][1])
    else:
        double_sushi.append((sushi_list[i]))
manzoku = 0
sushi_list = sushi_list[K:N]
for i in range(len(eat_sushi)):
    manzoku += eat_sushi[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku
flg = 0
while double_sushi:
    if(flg == 1):
        break
    else:
        tmp1 = double_sushi.pop()
        max_val = [0,0]
        for i in range(len(sushi_list)):
            if(sushi_list[i][1] not in eaten_sushi):
                tmp = sushi_list[i]
                max_val = sushi_list[i]
                break
            flg = 1
    tmp2 = max_val
    if(tmp2[0] == 0 and tmp2[1] == 0):
        break
    eaten_sushi.add(tmp2[1])
    tmp_manzoku = tmp_manzoku + tmp2[0] - tmp1[0] + len(eaten_sushi) ** 2 - (len(eaten_sushi) - 1) ** 2
    if (tmp_manzoku > manzoku):
        manzoku = tmp_manzoku
print(manzoku)