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 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 = tuple(l)
ans = []
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])
        ans.append(sushi_list[i])
    else:
        ans.append(sushi_list[i])
        double_sushi.append((sushi_list[i]))
manzoku = 0
for i in range(len(ans)):
    manzoku += ans[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku 
while double_sushi:
    tmp1 = double_sushi.pop()
    tmp2 = maxsushi(sushi_list,eaten_sushi)
    if(tmp2[0] == 0 and tmp2[1] == 1):
        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]
            if(tmp[0]) > max_val[0]:
                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 = tuple(l)
ans = []
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])
        ans.append(sushi_list[i])
    else:
        ans.append(sushi_list[i])
        double_sushi.append((sushi_list[i]))
manzoku = 0
for i in range(len(ans)):
    manzoku += ans[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku 
while double_sushi:
    tmp1 = double_sushi.pop()
    tmp2 = maxsushi(sushi_list,eaten_sushi)
    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]
            if(tmp[0]) > max_val[0]:
                max_val = sushi_list[i]
    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 = tuple(l)
ans = []
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])
        ans.append(sushi_list[i])
    else:
        ans.append(sushi_list[i])
        double_sushi.append((sushi_list[i]))
manzoku = 0
for i in range(len(ans)):
    manzoku += ans[i][0]
manzoku = manzoku + (len(eaten_sushi)) ** 2
tmp_manzoku = manzoku 
while double_sushi:
    tmp1 = double_sushi.pop()
    tmp2 = maxsushi(sushi_list,eaten_sushi)
    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:
# /usr/bin/python
# -*- coding: utf-8 -*-
import sys


n,k = map(int, input().split())
td_mx = {}
td = []
for _ in range(n):
  t,d = map(int, input().split())
  if t not in td_mx:
    td_mx[t] = d
  elif td_mx[t] < d:
    td.append(td_mx[t])
    td_mx[t] = d
  else:
    td.append(d)
td = sorted(td, key=lambda x:-x)
td_mx = sorted(td_mx.values(), key=lambda x:-x)

# td,td_mx の累積和を求める
csum = [0] * (len(td)+1)
csum_mx = [0] * (len(td_mx)+1)
for i in range(len(td)):
  csum[i+1] = csum[i] + td[i]
for i in range(len(td_mx)):
  csum_mx[i+1] = csum_mx[i] + td_mx[i]


ans = 0
for i in range(1,k+1):
  tmp = csum_mx[min(i,len(csum_mx)-1)] +  csum[min(k-i,len(csum)-1)] + i*i
  ans = max(ans,tmp)
print(ans)

Submission status: Wrong Answer

Solution code:
# /usr/bin/python
# -*- coding: utf-8 -*-
import sys


n,k = map(int, sys.stdin.readline().rstrip().split())
td = []
td_mx = {}
for i in range(n):
  t,d = map(int, sys.stdin.readline().rstrip().split())
  if t not in td_mx:
    td_mx[t] = d
  elif td_mx[t] < d:
    td.append(td_mx[t])
    td_mx[t] = d
  else:
    td.append(d)

td = sorted(td, key=lambda x:-x)
td_mx = sorted(td_mx.values(), key=lambda x:-x)
csum = [0]
for i in range(len(td)):
  if i == 0:
    csum.append(td[i])
  else:
    csum.append(csum[-1] + td[i])
csum_mx = [0]
for i in range(len(td_mx)):
  if i == 0:
    csum_mx.append(td_mx[i])
  else:
    csum_mx.append(csum_mx[-1] + td_mx[i])

ans = 0
for i in range(1,k+1):
  tmp = csum[min(k-i, len(csum)-1)]+csum_mx[min(i, len(csum_mx)-1)]+i**2
  ans = max(ans,tmp)
print(ans)

Submission status: Wrong Answer

Solution code:
import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(set(t1[0:K])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2,key=lambda x: x[1], reverse=True)
t2,d2 = zip(*td2)
s = set()
ans2 = 0
counter = 0
for i in range(N):
    if(len(s) == K):
        break
    if(t2[i] not in s):
        ans2 +=d2[i]  
        s.add(t2[i])
ans2 = ans2 + K ** 2
print(max(ans1,ans2))

Submission status: Wrong Answer

Solution code:
import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(set(t1[0:K-1])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2,key=lambda x: x[1], reverse=True)
t2,d2 = zip(*td2)
s = set()
ans2 = 0
counter = 0
for i in range(N):
    if(len(s) == K):
        break
    if(t2[i] not in s):
        ans2 +=d2[i]  
        s.add(t2[i])
ans2 = ans2 + K ** 2
print(max(ans1,ans2))

Submission status: Wrong Answer

Solution code:
import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(set(t1[0:K-1])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2,key=lambda x: x[1], reverse=True)
t2,d2 = zip(*td2)
s = set()
ans2 = 0
counter = 0
for i in range(N):
    if(len(s) == K):
        break
    if(t2[i] not in s):
        ans2 +=d2[i]  
        s.add(t2[i])
print(ans2)
ans2 = ans2 + K ** 2
print(max(ans1,ans2))

Submission status: Wrong Answer

Solution code:
import collections
N,K = map(int,input().split())
t = []
d = []
for _ in range(N):
    i,j = map(int,input().split())
    t.append(i)
    d.append(j)
td1 = list(zip(t,d))
td1 = sorted(td1,key=lambda x: x[1], reverse=True)
t1,d1 = zip(*td1)
ans1 = 0
for i in range(K):
    ans1 += d1[i]
ans1 = ans1 + len(collections.Counter(t1[0:K])) ** 2

td2 = list(zip(t,d))
td2 = sorted(td2)
ans2 = 0
s = set()
for i in range(K):
    if(td2[i][0] not in s):
        ans2 +=td2[i][1]
ans2 = ans2 + (K-1) ** 2
print(max(ans1,ans2))

Submission status: Wrong Answer

Solution code:
from collections import defaultdict


def main():
    """
    N個の寿司があり、それぞれの寿司には
    「ネタ」ti
    「おいしさ」di
    のパラメータが設定されています。 
    あなたはこの N個の寿司の中から K個を選んで食べようとしています。
    この時のあなたの「満足ポイント」は、以下のようにして計算されます。

    「満足ポイント」は、「おいしさ基礎ポイント」と、「種類ボーナスポイント」の和である。
    「おいしさ基礎ポイント」は、食べた寿司の「おいしさ」の総和である。
    「種類ボーナスポイント」は、食べた寿司の「ネタ」の種類数を xとしたとき、x ∗ x である。

    あなたは、「満足ポイント」をできるだけ大きくしたいです。
    この時の「満足ポイント」の値を求めてください。

    1 ≦ K  ≦ N ≦ 10^5
    1 ≦ ti ≦ N
    1 ≦ di ≦10^9

    nCk の組合せ: 10^5 C 10^5/2 => TLE
    DP ?
        1-Nまで, 食べる食べないによるスコア
        食べたとして種類数xとして種類は別途持っておく？
        →二次元配列は無理っぽそう, 10^5 * 10^5 = 10^10
        MAX時の種類だけ持って、１次元にする
        K までになっているかの判定
        →Nとは別に持たないといけない
    """
    N, K = map(int, input().split())
    td = [
        list(map(int, input().split()))
        for _ in range(N)
    ]
    ans = editorial(N, K, td)
    print(ans)


def editorial(N, K, td):
    td = sorted(td, key=lambda x: -x[1])
    c = defaultdict(int)
    d_sum = 0
    # 種類数が増えないので、不要と判断して良い基礎ポイント
    td_not_first = []
    for t, d in td[:K]:
        if c[t] > 0:
            td_not_first.append(d)
        c[t] += 1
        d_sum += d

    t_sum = len(c)
    ans = d_sum + t_sum ** 2

    for t, d in td[K:]:
        # 更新の余地なし
        if not td_not_first:
            return ans
        # すでに存在するなら、基礎ポイント降順のため不要
        if c[t] > 0:
            continue

        d_sum += -td_not_first.pop() + d
        t_sum += 1
        ans = max(ans, d_sum + t_sum ** 2)

    return ans
    # ans = editorial_WA_TLE(N, K, td)


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
def main():
    """
    N個の寿司があり、それぞれの寿司には
    「ネタ」ti
    「おいしさ」di
    のパラメータが設定されています。 
    あなたはこの N個の寿司の中から K個を選んで食べようとしています。
    この時のあなたの「満足ポイント」は、以下のようにして計算されます。

    「満足ポイント」は、「おいしさ基礎ポイント」と、「種類ボーナスポイント」の和である。
    「おいしさ基礎ポイント」は、食べた寿司の「おいしさ」の総和である。
    「種類ボーナスポイント」は、食べた寿司の「ネタ」の種類数を xとしたとき、x ∗ x である。

    あなたは、「満足ポイント」をできるだけ大きくしたいです。
    この時の「満足ポイント」の値を求めてください。

    1 ≦ K  ≦ N ≦ 10^5
    1 ≦ ti ≦ N
    1 ≦ di ≦10^9

    nCk の組合せ: 10^5 C 10^5/2 => TLE
    DP ?
        1-Nまで, 食べる食べないによるスコア
        食べたとして種類数xとして種類は別途持っておく？
        →二次元配列は無理っぽそう, 10^5 * 10^5 = 10^10
        MAX時の種類だけ持って、１次元にする
        K までになっているかの判定
        →Nとは別に持たないといけない
    """
    N, K = map(int, input().split())
    td = [
        list(map(int, input().split()))
        for _ in range(N)
    ]
    ans = editorial(N, K, td)
    print(ans)


def editorial(N, K, td):
    td = sorted(td, key=lambda x: -x[1])

    from collections import defaultdict
    c = defaultdict(int)
    s = defaultdict(list)
    d_sum = 0
    for t, d in td[:K]:
        d_sum += d
        s[t].append(d)
        c[t] += 1
    t_sum = len(c)
    ans = d_sum + t_sum ** 2

    # print(ans)
    # print(td[:K])
    # print(td[K:])
    for t, d in td[K:]:
        # すでに存在するなら、基礎ポイント降順のため不要
        if c[t] > 0:
            continue

        for t2, count in c.items():
            # 基礎が大きく、1つしかないならば必要なので残す
            if count <= 1:
                continue

            # print((t, d), t2, s[t], s[t2], c[t], c[t2], d_sum, t_sum, sep="\t")
            # 消しても種類数が増えないものを除外して、種類数が増えるように調整
            ex = s[t2].pop()
            s[t].append(d)

            c[t] += 1
            c[t2] -= 1
            d_sum += -ex + d
            t_sum += 1
            # print((t, d), t2, s[t], s[t2], c[t], c[t2], d_sum, t_sum, sep="\t")

            tmp = d_sum + t_sum ** 2
            ans = max(ans, tmp)
            break

    return ans


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
from collections import Counter
n, k = list(map(int, input().split(' ')))
t = []
d = []
for i in range(n):
    a, b = list(map(int, input().split(' ')))
    t.append(a)
    d.append(b)
    
def point(x):
    oishisa = sum(d[i] for i in x)
    nt = len(set(t[i] for i in x))
    return oishisa + nt * nt

c = list(range(k))
ct = Counter(t[i] for i in c)
for i in range(k, n):
    for j in range(k):
        if d[i] > d[c[j]]:# or t[i] != t[c[j]]:
            c[j] = i
print(point(c))

Submission status: Wrong Answer

Solution code:
N, K = map(int,input().split())
T = {}
for i in range(N):
    t, d = map(int, input().split())
    if t not in T.keys():
        T[t] = []
    T[t].append(d)

for i in range(N):
    T[t] = sorted(T[t], reverse=True)

T1 = []
T2 = []
for t in T.keys():
    T1.append(T[t][0])
    T2 += T[t][1:]

T1 = sorted(T1, reverse=True)
T2 = sorted(T2, reverse=True)
max_sc = 0
max_t = min(len(T1),K)
for i in range(1,max_t+1):
    sc = sum(T1[:i])
    sc += sum(T2[:K-i])
    sc += i * i
    max_sc = max(max_sc, sc)

print(max_sc)

Submission status: Wrong Answer

Solution code:
# N, K = map(int, input().split())
# sushi = []
# for i in N:
#     s = list(map(int, input().split()))
#     sushi.append(s)

# sushi = [
#     [1, 1],
#     [2, 1],
#     [3, 1],
#     [4, 6],
#     [4, 5],
#     [4, 5],
#     [4, 5],]
# N = len(sushi)
# K = 4
sushi = [
    [1, 9],
    [1, 7],
    [2, 6],
    [2, 5],
    [3, 1]]
N = len(sushi)
K = 3

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

def find_next(start, curr, point, tlist, tcount):
    if start == N:
        return [point, curr]
    max_c = curr
    max_p = point
    for s in sushi[start:]:
        if tlist[s[0]] > 0:
            continue
        for ri in range(len(curr) - 1, -1, -1):
            r = curr[ri]
            if tlist[r[0]] >= 2:
                c2 = [c for c in curr]
                t2 = [t for t in tlist]
                del c2[ri]
                c2.append(s)
                t2[r[0]] -= 1
                t2[s[0]] = 1
                p2 = point - r[1] + s[1] + (tcount + 1) * (tcount + 1) - tcount * tcount
                np, pc = find_next(start + 1, c2, p2, t2, tcount + 1)

                if np > max_p:
                    max_p = np
                    max_c = pc
                break
    return [max_p, max_c]

result = sushi[0: K]
type_count = [0] * (N + 1)

for r in result:
    type_count[r[0]] += 1
type_num = 0
for t in type_count:
    if t > 0:
        type_num += 1
point = type_num * type_num
for r in result:
    point += r[1]
ans_p, ans_c = find_next(K, result, point, type_count, type_num)

print(ans_p)

Submission status: Wrong Answer

Solution code:
import collections
import itertools
import numpy as np
import sys
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(sys.stdin.readline())
def S(): return input()

N,K = IL()
l = [[] for i in range(N+1)]
for i in range(N):
    a,b = IL()
    l[a].append(b)

l.sort(reverse=True)

one = []
zero = []
for i in range(N+1):
    if len(l[i]) == 0:
        continue
    else:
        one.append(l[i].pop(0))
for i in l:
    zero.extend(i)

one.sort(reverse=True)
zero.sort(reverse=True)

onel = len(one)
zerol = len(zero)

if K-1 <= zerol:
    start = 0
else:
    start = (K-1)-zerol
end = min(onel,K)

ans = 0
tmp = 0
for i in range(start,end):
    if i == start:
        ans = sum(one[:i+1])+sum(zero[:K-1-i])+(i+1)**2
        tmp = ans
    else:
        tmp += one[i]-zero[K-1-i]-i**2+(i+1)**2
        ans = max(ans,tmp)   
print(ans)

Submission status: Wrong Answer

Solution code:
import collections
import itertools
import numpy as np
import sys
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(sys.stdin.readline())
def S(): return input()

N,K = IL()
l = [[] for i in range(N+1)]
for i in range(N):
    a,b = IL()
    l[a].append(b)

l.sort(reverse=True)

one = []
zero = []
for i in range(N+1):
    if len(l[i]) == 0:
        continue
    else:
        one.append(l[i].pop(0))
for i in l:
    zero.extend(i)

one.sort(reverse=True)
zero.sort(reverse=True)

onel = len(one)
zerol = len(zero)

start = K-min(zerol,K-1)-1
end = min(onel,K)

ans = 0
tmp = 0
for i in range(start,end):
    if i == start:
        ans = sum(one[:i+1])+sum(zero[:K-1-i])+(i+1)**2
        tmp = ans
    else:
        tmp += one[i]-zero[K-1-i]-i**2+(i+1)**2
        ans = max(ans,tmp)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
td = []
sushi_max = [-1] * N
for i in range(N):
    t,d = map(int,input().split(" "))
    td.append((d,t))
    sushi_max[t-1] = max(sushi_max[t-1],d)

td.sort(reverse=True,key=lambda x:x[0])
cnt = [0] * N
que = []
ans = 0
for i in range(K):
    cnt[td[i][1]-1] += 1
    sushi_max[td[i][1]-1] = -1
    pq.heappush(que,td[i])
    ans += td[i][0]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz*sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(que) > 0:
    minValue,minIdx = pq.heappop(que)
    if cnt[minIdx-1] == 1:
        continue
    cur -= minValue
    cur -= (sz*sz)
    cnt[minIdx-1] -= 1
    cnt[sushi_max[idx][1]-1] += 1
    cur += (sushi_max[idx][0] + (sz+1)*(sz+1))
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
td = []
sushi_max = [-1] * N
for i in range(N):
    t,d = map(int,input().split(" "))
    td.append((d,t))
    sushi_max[t-1] = max(sushi_max[t-1],d)

td.sort(reverse=True,key=lambda x:x[0])
cnt = [0] * N
que = []
ans = 0
for i in range(K):
    cnt[td[i][1]-1] += 1
    sushi_max[td[i][1]-1] = -1
    pq.heappush(que,td[i])
    ans += td[i][0]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz*sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(que) > 0:
    minValue,minIdx = pq.heappop(que)
    if cnt[minIdx-1] == 0:
        continue
    cur -= minValue
    cur -= (sz*sz)
    cnt[minIdx-1] -= 1
    cnt[sushi_max[idx][1]-1] += 1
    cur += (sushi_max[idx][0] + (sz+1)*(sz+1))
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
td = []
sushi_max = [-1] * N
for i in range(N):
    t,d = map(int,input().split(" "))
    td.append((d,t))
    sushi_max[t-1] = max(sushi_max[t-1],d)

td.sort(reverse=True,key=lambda x:x[0])
cnt = [0] * N
que = []
ans = 0
for i in range(K):
    cnt[td[i][1]-1] += 1
    sushi_max[td[i][1]-1] = -1
    pq.heappush(que,td[i])
    ans += td[i][0]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz*sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(que) > 0:
    minValue,minIdx = pq.heappop(que)
    cur -= minValue
    cur -= (sz*sz)
    cnt[minIdx-1] -= 1
    if cnt[minIdx-1] == 0 or sushi_max[idx][0] == -1:
        break
    cnt[sushi_max[idx][1]-1] += 1
    cur += sushi_max[idx][0]
    cur += (sz+1)*(sz+1)
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
import heapq as pq
N,K = map(int,input().split(" "))
td = []
sushi_max = [-1] * N
for i in range(N):
    t,d = map(int,input().split(" "))
    td.append((d,t))
    sushi_max[t-1] = max(sushi_max[t-1],d)

td.sort(reverse=True,key=lambda x:x[0])
cnt = [0] * N
que = []
ans = 0
for i in range(K):
    cnt[td[i][1]-1] += 1
    sushi_max[td[i][1]-1] = -1
    pq.heappush(que,td[i])
    ans += td[i][0]

sz = len(list(filter(lambda x: x > 0,cnt)))
ans += (sz*sz)

sushi_max = sorted([ (val,idx) for idx,val in enumerate(sushi_max) ],reverse=True,key=lambda x:x[0])
idx = 0
cur = ans
while len(que) > 0:
    minValue,minIdx = pq.heappop(que)
    cur -= minValue
    cur -= (sz*sz)
    cnt[minIdx-1] -= 1
    if cnt[minIdx-1] == 0:
        break
    cnt[sushi_max[idx][1]-1] += 1
    cur += sushi_max[idx][0]
    cur += (sz+1)*(sz+1)
    sz += 1
    idx += 1
    ans = max(cur,ans)
print(ans)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(100000)


def read_int():
    return int(input())


def read_int_n():
    return list(map(int, input().split()))


def read_float():
    return float(input())


def read_float_n():
    return list(map(float, input().split()))


def read_str():
    return input().strip()


def read_str_n():
    return list(map(str, input().split()))


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap


@mt
def slv(N, K, TD):
    TD.sort(key=lambda x: x[1], reverse=True)
    # for r in TD:
    #     print(r)

    c = Counter()
    ans = 0
    rc = []
    for i in range(K):
        t, k = TD[i]
        c[t] += 1
        ans += k
        if c[t] > 1:
            rc.append((t, k))
    
    mi = i
    while True:
        if len(rc) == 0:
            break
        rt, rk = rc.pop()
        cc = len(c)
        d = (cc+1)**2 - cc**2
        for i in range(mi+1, len(TD)):
            if (TD[i][0] not in c) and (d > rk - TD[i][1]):
                c[TD[i][0]] += 1
                c[rt] -= 1
                ans -= rk
                ans += TD[i][1]
                mi = i
                break

        else:
            break
    # print(c)
    return ans + len(c)**2


def main():
    # N = read_int()
    N, K = read_int_n()
    TD = [read_int_n() for _ in range(N)]
    print(slv(N, K, TD))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(100000)


def read_int():
    return int(input())


def read_int_n():
    return list(map(int, input().split()))


def read_float():
    return float(input())


def read_float_n():
    return list(map(float, input().split()))


def read_str():
    return input().strip()


def read_str_n():
    return list(map(str, input().split()))


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap


@mt
def slv(N, K, TD):
    TD.sort(key=lambda x: x[1], reverse=True)
    # for r in TD:
    #     print(r)

    c = Counter()
    ans = 0
    for i in range(K):
        t, k = TD[i]
        c[t] += 1
        ans += k
    
    mi = i
    while True:
        cc = len(c)
        d = (cc+1)**2 - cc**2
        for i in range(mi+1, len(TD)):
            if (TD[i][0] not in c) and (d > TD[mi][1] - TD[i][1]):
                # print(mi, i, TD[mi], TD[i])
                c[TD[i][0]] += 1
                c[TD[mi][0]] -= 1
                ans -= TD[mi][1]
                ans += TD[i][1]
                mi = i
                break

        else:
            break
    # print(c)
    return ans + len(c)**2


def main():
    # N = read_int()
    N, K = read_int_n()
    TD = [read_int_n() for _ in range(N)]
    print(slv(N, K, TD))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
def various_sushi(N: int, K: int, TD: list) -> int:
    _TD = sorted(TD, key=lambda x: x[1], reverse=True)
    tops, others = [], []
    tops_sum, others_sum = 0, 0
    kind = set()
    for t, d in _TD:
        if t not in kind:
            tops_sum += d
            tops.append(tops_sum)

            kind.add(t)
        else:
            others_sum += d
            others.append(others_sum)

    max_d = 0
    types = len(kind)
    for p in range(types):
        if K < p:
            break
        delicious = tops[p] + (p + 1) * (p + 1)
        index = (K - 1) - (p + 1)
        if 0 <= index and index < N - types:
            delicious += others[index]

        max_d = max(max_d, delicious)

    return max_d


if __name__ == "__main__":
    N, K = [int(s) for s in input().split()]
    TD = []
    for _ in range(N):
        t, d = [int(s) for s in input().split()]
        TD.append((t, d))

    ans = various_sushi(N, K, TD)
    print(ans)

Submission status: Wrong Answer

Solution code:
def various_sushi(N: int, K: int, TD: list) -> int:
    _TD = sorted(TD, key=lambda x: x[1], reverse=True)
    tops, others = [], []
    tops_sum, others_sum = 0, 0
    kind = set()
    for t, d in _TD:
        if t not in kind:
            tops_sum += d
            tops.append(tops_sum)

            kind.add(t)
        else:
            others_sum += d
            others.append(others_sum)

    max_d = 0
    for p in range(len(kind)):
        if K < p:
            break
        delicious = tops[p] + (p + 1) * (p + 1)
        if 0 <= (K - 1) - (p + 1) and (K - 1) - (p + 1) < N - len(kind):
            delicious += others[(K - 1) - (p + 1)]

        max_d = max(max_d, delicious)

    return max_d


if __name__ == "__main__":
    N, K = [int(s) for s in input().split()]
    TD = []
    for _ in range(N):
        t, d = [int(s) for s in input().split()]
        TD.append((t, d))

    ans = various_sushi(N, K, TD)
    print(ans)

Submission status: Wrong Answer

Solution code:
def various_sushi(N: int, K: int, TD: list) -> int:
    _TD = sorted(TD, key=lambda x: x[1], reverse=True)
    most_delicious_by_type = []
    others = []
    first = True
    used_types = []  # most_delicious_by_type に入っている type
    types = 0
    for t, d in _TD:
        if t not in used_types:
            if len(most_delicious_by_type) == 0:
                most_delicious_by_type.append(d)
            else:
                most_delicious_by_type.append(most_delicious_by_type[-1] + d)

            used_types.append(t)
            types += 1
        else:
            if first:
                others.append(d)
                first = False
            else:
                others.append(others[-1] + d)

    most_delicious = 0
    for p in range(0, types):
        if K < p:
            break
        delicious = most_delicious_by_type[p] + (p + 1) * (p + 1)
        if 0 <= (K - 1) - (p + 1) and (K - 1) - (p + 1) < N - types:
            delicious += others[(K - 1) - (p + 1)]

        most_delicious = max(most_delicious, delicious)

    return most_delicious


if __name__ == "__main__":
    N, K = [int(s) for s in input().split()]
    TD = []
    for _ in range(N):
        t, d = [int(s) for s in input().split()]
        TD.append((t, d))

    ans = various_sushi(N, K, TD)
    print(ans)

Submission status: Wrong Answer

Solution code:
from collections import deque
n, k = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(n)]
res = 0
td.sort(key = lambda x:-x[1])
kind = set()
queue = deque()
for i in range(k):
    if td[i][0] in kind:
        queue.append(td[i][1])
    else:
        kind = kind | {td[i][0]}
    res += td[i][1]
ans1 = res + len(kind)**2
point = 0
count = len(queue)
for i in range(k, n):
    if count == 0:
        break
    if not td[i][0] in kind:
        kind = kind | {td[i][0]}
        point += td[i][1]
        count -= 1
res = max(ans1, res - sum(queue) + point + len(kind)**2)
print(res)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])

ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans - q.pop() - len(s) ** 2 + TD[k1][1] + (len(s)+1) ** 2
        ans = max(ans, ans1)
        s.add(TD[k1][0])
    k1 += 1

print(ans)

Submission status: Wrong Answer

Solution code:
N, K = map(int, input().split())
TD = sorted([list(map(int, input().split())) for i in range(N)], key=lambda x: -x[1])
q = []
ans = 0
s = set()
for i in range(K):
    ans += TD[i][1]
    if TD[i][0] in s:
        q.append(TD[i][1])
    s.add(TD[i][0])

ans += len(s) ** 2
k1 = K
ans1 = ans
while k1 < N and q:
    if TD[k1][0] not in s:
        ans1 = ans - q.pop() - len(s) ** 2 + TD[k1][1] + (len(s)+1) ** 2
        ans = max(ans, ans1)
        if ans == ans1:
            s.add(TD[k1][0])
    k1 += 1
    
print(ans)

Submission status: Wrong Answer

Solution code:
N,K = (int(x) for x in input().split())
td_arr = []
for _ in range(N):
    td = tuple(int(x) for x in input().split())
    td_arr.append(td)

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

n_dm_set = set()
small_arr = []
d_dm = 0
for t,d in td_arr[:K]:
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

prev_t = -1
other_arr = []
for t,d in sorted(td_arr, key=lambda x:x[0]):
    if t == prev_t:
        continue
    if t in n_dm_set:
        continue
    prev_t = t
    other_arr.append((t,d))

d_x = d_dm
nt_x = len(n_dm_set)
for t,d in other_arr:
    if len(small_arr) == 0:
        break
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    nt_x += 1
    candidate_dict[nt_x] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)