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: Runtime Error

Solution code:
import math,string,itertools,fractions,heapq,collections,re,array,bisect,sys,random,time,copy,functools

sys.setrecursionlimit(10**7)
inf = 10**20
eps = 1.0 / 10**10
mod = 10**9+7
dd = [(-1,0),(0,1),(1,0),(0,-1)]
ddn = [(-1,0),(-1,1),(0,1),(1,1),(1,0),(1,-1),(0,-1),(-1,-1)]

def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LI_(): return [int(x)-1 for x in sys.stdin.readline().split()]
def LF(): return [float(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def I(): return int(sys.stdin.readline())
def F(): return float(sys.stdin.readline())
def S(): return input()
def pf(s): return print(s, flush=True)


def main():
    n,k = LI()
    dt = sorted([LI()[::-1] for _ in range(n)], reverse=True)

    u = set()
    s = 0
    o = []
    for d,t in dt[:k]:
        if t in u:
            o.append(d)
        else:
            u.add(t)
        s += d
    r = s + len(u) ** 2
    oi = len(o) - 1
    for d,t in dt[k:]:
        if t in u:
            continue
        s -= o[oi]
        oi -= 1
        s += d
        u.add(t)
        tr = s + len(u) ** 2
        if r < tr:
            r = tr
        if oi < 0:
            break

    return r



print(main())

Submission status: Runtime Error

Solution code:
def solve():
    n, k = map(int, input().split(' '))
    sushis = []
    for _ in range(n):
        kind, dp = map(int, input().split(' '))
        sushis.append((kind, dp))
    sushis.sort(key=lambda tp: tp[1], reverse=False)
    selected = []
    for i in range(k):
        selected.append(sushis.pop())

    max_point = cal_sp(selected)
    for i in range(num_kind(selected), k):
        # cal dp
        delete_same_sushi(selected)
        add_diff_sushi(selected, sushis)
        point = cal_sp(selected)
        if max_point < point:
            max_point = point
    print(max_point)

def cal_sp(selected): # satisfaction point
    kind = num_kind(selected)
    point = kind ** 2 + sum(map(lambda x: x[1], selected))
    return point

def add_diff_sushi(selected, sushis):
    while True:
        sushi = sushis.pop()
        if sushi[0] in map(lambda x: x[0], selected):
            continue
        else:
            selected.append(sushi)
            break

def delete_same_sushi(selected):
    delete_index = -1
    while True:
        last = selected[delete_index]
        if list(map(lambda x: x[0], selected)).count(last[0]) == 1:
            delete_index -= 1
        else:
            break
    selected.pop(delete_index)


def num_kind(sushis):
    kinds = set() 
    for sushi in sushis:
        kinds.add(sushi[0])
    return len(kinds)


if __name__ == "__main__":
    solve()

Submission status: Runtime Error

Solution code:
def solve():
    n, k = map(int, input().split(' '))
    sushis = []
    for _ in range(n):
        kind, dp = map(int, input().split(' '))
        sushis.append((kind, dp))
    sushis.sort(key=lambda tp: tp[1], reverse=False)
    selected = []
    for i in range(k):
        selected.append(sushis.pop())

    points = []
    points.append(cal_sp(selected))
    for i in range(num_kind(selected), k):
        # cal dp
        delete_same_sushi(selected)
        add_diff_sushi(selected, sushis)
        points.append(cal_sp(selected))
    print(max(points))

def cal_sp(selected): # satisfaction point
    kind = num_kind(selected)
    point = kind ** 2 + sum(map(lambda x: x[1], selected))
    return point

def add_diff_sushi(selected, sushis):
    while True:
        sushi = sushis.pop()
        if sushi[0] in map(lambda x: x[0], selected):
            continue
        else:
            selected.append(sushi)
            break

def delete_same_sushi(selected):
    delete_index = -1
    while True:
        last = selected[delete_index]
        if list(map(lambda x: x[0], selected)).count(last[0]) == 1:
            delete_index -= 1
        else:
            break
    selected.pop(delete_index)


def num_kind(sushis):
    kinds = set() 
    for sushi in sushis:
        kinds.add(sushi[0])
    return len(kinds)


if __name__ == "__main__":
    solve()

Submission status: Runtime Error

Solution code:
def solve():
    n, k = map(int, input().split(' '))
    sushis = []
    for _ in range(n):
        kind, dp = map(int, input().split(' '))
        sushis.append((kind, dp))
    sushis.sort(key=lambda tp: tp[1], reverse=False)
    selected = []
    for i in range(k):
        selected.append(sushis.pop())

    points = []
    points.append(cal_sp(selected))
    for k in range(num_kind(selected), kind):
        # cal dp
        delete_same_sushi(selected)
        add_diff_sushi(selected, sushis)
        points.append(cal_sp(selected))
    print(max(points))

def cal_sp(selected): # satisfaction point
    kind = num_kind(selected)
    point = kind ** 2 + sum(map(lambda x: x[1], selected))
    return point

def add_diff_sushi(selected, sushis):
    while True:
        sushi = sushis.pop()
        if sushi[0] in map(lambda x: x[0], selected):
            continue
        else:
            selected.append(sushi)
            break

def delete_same_sushi(selected):
    delete_index = -1
    while True:
        last = selected[delete_index]
        if list(map(lambda x: x[0], selected)).count(last[0]) == 1:
            delete_index -= 1
        else:
            break
    selected.pop(delete_index)


def num_kind(sushis):
    kinds = set() 
    for sushi in sushis:
        kinds.add(sushi[0])
    return len(kinds)


if __name__ == "__main__":
    solve()

Submission status: Runtime Error

Solution code:
def d_various_sushi(N, K, Param):
    import heapq

    # 寿司をおいしさの値が大きなものからK個取る
    sushi_parameter = sorted(Param, key=lambda x: x[1], reverse=True)
    sushi_invariably_chosen = []  # ネタが被っていない場合、貪欲に取っていい
    candidate_reject = []  # ネタの種類数を増やす時に取り除く候補(priority queueで管理)
    chosen_topping = set()  # 取った寿司ネタの種類
    for i in range(K):
        t, d = sushi_parameter[i]
        if t not in chosen_topping:
            sushi_invariably_chosen.append(d)
            chosen_topping.add(t)
        else:
            heapq.heappush(candidate_reject, d)

    # 暫定的な最大値を計算しておく
    delicious_sum1 = sum(sushi_invariably_chosen)
    delicious_sum2 = sum(candidate_reject)
    topping_variety = len(chosen_topping)
    ans = delicious_sum1 + delicious_sum2 + topping_variety**2
    if len(candidate_reject) == 0:
        return ans  # 除く寿司の候補がないので、満足度を大きくしようがない

    for i in range(K, N):
        t, d = sushi_parameter[i]
        if t in chosen_topping:
            # ネタが被った上、おいしさがこれまで取った寿司より小さな寿司と交換する理由はない
            continue

        # 取り除いてもネタの種類数が減らない寿司の中で最もおいしさの小さいものを取り除く
        delicious_sum2 -= heapq.heappop(candidate_reject)
        # これまでに取っている寿司にないネタを持ち、最もおいしさの大きい寿司を取る
        chosen_topping.add(t)
        delicious_sum1 += d
        topping_variety += 1

        ans = max(ans, delicious_sum1 + delicious_sum2 + topping_variety**2)
    return ans

N, K = [int(i) for i in input().split()]
Param = [[int(i) for i in input().split()] for j in range(N)]
print(d_various_sushi(N, K, Param))

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
umai = {}
mazui = []
for _ in range(N):
    t, d = map(int, input().split())
    if t not in umai:
        umai[t] = d
    else:
        tmp = umai[t]
        if tmp < d:
            umai[t] = d
            mazui.append((t, tmp))
        else:
            mazui.append((t, d))

umai = sorted(umai.items(), key=lambda x: -x[1])
mazui = sorted(mazui, key=lambda x: -x[1])


umai_len = len(umai)
mazui_len = len(mazui)

acc_umai = [0 for _ in range(umai_len)]
acc_umai[0] = umai[0][1]
for i in range(umai_len - 1):
    acc_umai[i + 1] = acc_umai[i] + umai[i + 1][1]

acc_mazui = [0 for _ in range(mazui_len)]
acc_mazui[0] = mazui[0][1]
for i in range(mazui_len - 1):
    acc_mazui[i + 1] = acc_mazui[i] + mazui[i + 1][1]
acc_mazui.insert(0, 0)

ans = []
for i in range(1, K + 1):
    if umai_len < i or mazui_len < K - i:
        continue
    ans.append(acc_umai[i - 1] + acc_mazui[K - i] + i ** 2)
print(max(ans))

Submission status: Runtime Error

Solution code:
N, K = [int(_) for _ in input().split()]
TD = [[int(_) for _ in input().split()] for i in range(N)]
 
# print(TD)
 
xs = sorted(TD, key=lambda x:-x[1])
 
#print(xs)
 
baset = 0
ts = set()
ds = []
 
for t, d in xs[:K]:
    baset += d
    if not t in ts:
        ts.add(t)
    else:
        ds.append(d)
result = baset + len(ts)**2
r = result

for t, d in xs[K:]:
    if not ts:
        break
    if t not in ts:
        r += 2*len(ts) + 1
        r += d - ds.pop()
        ts.add(t)
        result = max(r,result)

print(result)

Submission status: Runtime Error

Solution code:
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <cmath>
using namespace std;
using ll = long long;

struct P {
    bool operator<(const P& p) const {
        return d > p.d;
    }
    int t, d;
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n, k;
    cin >> n >> k;

    vector<P> p(n);
    int m = 0;
    for (int i = 0; i < n; i++) {
        int t, d;
        cin >> t >> d;
        t--;

        p[i] = { t, d };
        m = max(m, t);
    }
    sort(p.begin(), p.end());
    m++;

    vector<int> d(n);
    vector<char> a(m);
    int i0 = 0, i1 = 0;
    ll r = 0, s = 0;

    for (int i = 0; i < n; i++) {
        if (a[p[i].t] == 0) {
            a[p[i].t] = 1;
            if (i >= k) {
                if (i1 == 0) break;
                s -= d[--i1];
            }
            i0++;
            s += p[i].d;
        } else {
            if (i < k) s += d[i1++] = p[i].d;
        }
        r = max(r, s + (ll)i0 * i0);
    }

    cout << r << endl;

    return 0;
}

Submission status: Runtime Error

Solution code:
from collections import  deque
N,K = map(int, input().split())
neta_uma = []
for i in range(N):
    neta,uma = map(int, input().split())
    neta_uma.append((neta, uma))
neta_uma = sorted(neta_uma, key=lambda x:x[1], reverse=True)
uma_max = [0]*(N+1) # ネタごとのうまさの最大
for neta,uma in neta_uma:
    uma_max[neta] = max(uma_max[neta], uma)
# print(uma_max)

# うまさ全振り
chosen = neta_uma[:K]
used = set(n for n,u in chosen)
# print(chosen, used)
ans = sum(u for n,u in chosen) + len(used)**2

# 多様性を増やす方向
tmp = ans
for i in range(K,N):
    # print(neta_uma[i], ans, used, )
    if not chosen: break
    
    if neta_uma[i][0] not in used:
        for j in reversed(range(K)):
            neta,uma = chosen[j]
            if uma==uma_max[neta]: continue
            chosen.remove(chosen[j])
            tmp += neta_uma[i][1] - uma
            tmp += 2*len(used) +1     # (k+1)^2 - k^2
            used.add(neta_uma[i][0])
            ans = max(ans, tmp)
        
print(ans)

Submission status: Runtime Error

Solution code:
from math import pow
def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        # default
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    # optimzie at most K times
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: # no optimization occured
            break
        else:
            score = new_score
    return score
  
N, K = map(int, input().split())
records = [tuple(map(int, input().split())) for i in range(N)] 
score =cal_score(N, K, records)
print(int(score))

Submission status: Runtime Error

Solution code:
from math import pow

def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        # default
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    # optimzie at most K times
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: # no optimization occured
            break
        else:
            score = new_score
    return score
  
N, K = map(int, input().split())
records = [list(map(int, input().split())) for i in range(N)] 
score = cal_score(N, K, records)
print(score)

Submission status: Runtime Error

Solution code:
from math import pow






def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        # default
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    # optimzie at most K times
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: # no optimization occured
            break
        else:
            score = new_score
    return score
  
N, K = map(int, input().split())
records = [tuple(map(int, input().split())) for i in range(N)] 
score = cal_score(N, K, records)
print(score)

Submission status: Runtime Error

Solution code:
from math import pow



N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]


def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        # default
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    # optimzie at most K times
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: # no optimization occured
            break
        else:
            score = new_score
    return score
 
score = cal_score(N, K, records)
print(score)

Submission status: Runtime Error

Solution code:
from math import pow



N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]


def compute_score(target_collection):
    keys = {neta_point[0] for neta_point in target_collection}
    points = [neta_point[1] for neta_point in target_collection]
    cardinality = len(keys)
    score = pow(cardinality, 2) + sum(points)
    return score

# exchange one record and compare
def optimize(full_collection, target_collection, target_collection_len, idx):
    while True:
        ex_record = full_collection[idx]
        # check on the idx-th record
        idx += 1
        # bump up for return value
        if ex_record in target_collection:
            continue # cardinality score will not increase
        else:
            for i in range(target_collection_len-1):
                if target_collection[i][0] == target_collection[i-1][0]:
                    target_collection[i] = ex_record
                    new_score = compute_score(target_collection)
                    return new_score, target_collection, idx
        # default
        return 0, target_collection, idx



def cal_score(N, K, records):
    sorted_recs = sorted(records, reverse=True, key=lambda x: x[1])
    target_collection = sorted_recs[0:K]
    score =  compute_score(target_collection)

    # optimzie at most K times
    idx = K
    for i in range(K):
        new_score, new_target_collection, idx = optimize(sorted_recs, target_collection, K, idx)
        if new_score <= score: # no optimization occured
            break
        else:
            score = new_score
    return score
 
score =cal_score(N, K, record)
print(score)

Submission status: Runtime Error

Solution code:
from operator import itemgetter
from collections import Counter
import heapq
N, K = map(int, input().split())
td = [list(map(int, input().split())) for _ in range(N)]
td.sort(key=itemgetter(1, 0), reverse=True)
t, d = zip(*td)

cnt_all = Counter(t)
cnt_all = cnt_all.most_common()
kind = len(cnt_all)
tmp = 0
i = 0
while tmp < K:
    tmp += cnt_all[i][1]
    i += 1
init_kind = i
S = []
heapq.heapify(S)
cnt = Counter()
for i in range(K):
    cnt[t[i]] += 1
    heapq.heappush(S, (d[i], i, t[i]))
sums = [0]*(K+1)
sums[init_kind] = sum(list(zip(*S))[0])
for i in range(init_kind, kind):
    nS = []
    heapq.heapify(nS)
    keys = cnt.keys()
    remove_flag = False
    for _ in range(K):
        tmp = heapq.heappop(S)
        if cnt[tmp[2]] == 1:
            heapq.heappush(nS, tmp)
        else:
            if remove_flag:
                heapq.heappush(nS, tmp)
            else:
                remove_flag = True
                cnt[tmp[2]] -= 1
    for j in range(N):
        if t[j] not in keys:
            heapq.heappush(nS, (d[j], j, t[j]))
            cnt[t[j]] += 1
            break
    sums[i+1] = sum(list(zip(*nS))[0])
    S = nS
for i in range(len(sums)):
    sums[i] += i**2
print(max(sums))

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
from collections import deque
n,k = map(int, input().split())
td = []

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

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

used = set()
q = []
res = 0
kind = 0
for i in range(k):
    t,d = td[i]
    if t not in used:
        kind += 1
        used.add(t)
    else:
        q.append((t,d))
    res += d
res += kind**2
q = deque(sorted(q, key=lambda x:x[1]))

tmp = res
for i in range(k,n):
    t2,d2 = td[i]
    if t2 in used:
        continue
    t,d = q.popleft()

    r = tmp - d + d2
    r -= kind**2
    kind += 1
    r += kind**2
    res = max(res, r)
    tmp = r
    if len(q)==0:
        break

print(res)

Submission status: Runtime Error

Solution code:
N, K = map(int, input().split())
sushi = [tuple(map(int, input().split())) for i in range(N)]
point = [0 for i in range(K+1)]

sushi.sort(lambda x: x[1])

eaten_sushi = [sushi[i] for i in range(K)]
eaten_variation = [s[0] for s in eaten_sushi]
variation_number = len(set(eaten_variation))
base = sum([s[1] for s in eaten_sushi])
bonus = variation_number ** 2
point[variation_number] = base + bonus

for i in range(K, N):
    if sushi[i][0] in eaten_variation:
        continue

    for j in range(K-1, 0, -1):
        if eaten_variation.count(eaten_sushi[j][0]) == 1:
            continue
        eaten_sushi.pop(j)
        eaten_sushi.append(sushi[i])
        eaten_variation = [s[0] for s in eaten_sushi]
        variation_number = len(set(eaten_variation))
        base = sum([s[1] for s in eaten_sushi])
        bonus = variation_number ** 2
        point[variation_number] = base + bonus
        break

print(max(point))

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))
A.sort(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
len_x = len(X)
len_y = len(Y)
if len_y >= k:
    a = 0
    b = len_y - 1
    x = 0
    y = sum(Y[:k])
else:
    a = k - len_y
    b = len_y - 1
    x = sum(X[:k - len_y])
    y = sum(Y)
ans = x + y + len_y ** 2
t = ans
for i in range(min(len_x, len_y, k)):
    x += X[a]
    a += 1
    y -= Y[b]
    b -= 1
    len_y -= 1
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))
A.sort(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
len_x = len(X)
len_y = len(Y)
if len_y >= k:
    a = 0
    b = len_y - 1
    x = 0
    y = sum(Y[:k])
    N = min(k, len_x)
else:
    a = k - len_y
    b = len_y - 1
    x = sum(X[:k - len_y])
    y = sum(Y)
    N = min(k - len_y, len_x)
ans = x + y + len_y ** 2
t = ans
for i in range(N):
    x += X[a]
    a += 1
    y -= Y[b]
    b -= 1
    len_y -= 1
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))
A.sort(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
len_y = len(Y)
if len_y >= k:
    a = 0
    b = len_y - 1
    x = 0
    y = sum(Y[:k])
else:
    a = k - len_y
    b = len_y - 1
    x = sum(X[:k - len_y])
    y = sum(Y)
ans = x + y + len_y ** 2
t = ans
for i in range(min(len(X), len_y, k - len_y + 1, k - 1)):
    x += X[a]
    a += 1
    y -= Y[b]
    b -= 1
    len_y -= 1
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Runtime Error

Solution code:
def IL(): return list(map(int,input().split()))
def SL(): return input().split()
def I(): return int(input())
def S(): return input()
n,k=IL()
top=dict()
sub=[]
for i in range(n):
    t,d=IL()
    if t in top:
        sub.append(min(top[t],d))
        top[t]=max(top[t],d)
    else:
        top[t]=d

top=sorted(top.items(),key=lambda x:-x[1])
sub.sort(reverse=True)

rtop=[top[0][1]]
rsub=[sub[0]]
for i in range(1,len(top)):
    rtop.append(rtop[-1]+top[i][1])
for i in range(1,len(sub)):
    rsub.append(rsub[-1]+sub[i])

ans=0
for i in range(k):
    if i<len(rtop) and 0<=k-i-2<len(rsub):
        ans=max(ans,rtop[i]+rsub[k-i-2]+(i+1)**2)
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))
A.sort(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
len_y = len(Y)
if len_y >= k:
    a = 0
    b = len_y - 1
    x = 0
    y = sum(Y[:k])
else:
    a = k - len_y
    b = len_y - 1
    x = sum(X[:k - len_y])
    y = sum(Y)
ans = x + y + len_y ** 2
t = ans
for i in range(min(len(X), k - len_y + 1, k - 1)):
    x += X[a]
    a += 1
    y -= Y[b]
    b -= 1
    len_y -= 1
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Runtime Error

Solution code:
n, k = map(int, input().split())
A = []
for _ in range(n):
    t, d = map(int, input().split())
    A.append((d, t))
A.sort(reverse=True)
X = []
Y = []
st = set()
for d, t in A:
    if t in st:
        X.append(d)
    else:
        Y.append(d)
    st.add(t)
ans = 0
len_y = len(Y)
if len_y >= k:
    a = 0
    b = len_y - 1
    x = 0
    y = sum(Y[:k])
else:
    a = k - len_y
    b = len_y - 1
    x = sum(X[:k - len_y])
    y = sum(Y)
ans = x + y + len_y ** 2
t = ans
for i in range(min(len(X), k - 1)):
    x += X[a]
    a += 1
    y -= Y[b]
    b -= 1
    len_y -= 1
    t = x + y + len_y ** 2
    ans = max(ans, t)
print(ans)

Submission status: Runtime Error

Solution code:
import heapq as h;I=lambda:list(map(int,input().split()));n,k=I();z=[i*2+1for i in range(n)];s=sorted([I()for _ in[0]*n],key=lambda x:-x[1]);q=[];v=T=0;l=[0]*-~n
for a in s[:k]:
 if l[a[0]]:h.heappush(q,a[1])
 else:A+=z[v];v+=1
 l[a[0]]=1;A+=a[1]
T=A
for a in s[k:]:
 if not q:break
 if l[a[0]]^1:l[a[0]]=1;t=h.heappop(q);T=T-t+a[1]+z[v];v+=1;A=max(A,T)
print(A)

Submission status: Runtime Error

Solution code:
# coding:utf-8

import sys
import heapq
from collections import defaultdict, deque

INF = float('inf')
MOD = 10 ** 9 + 7

def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def II(): return int(sys.stdin.readline())
def SI(): return input()


n, k = LI()
sushi = [(d, t) for t, d in (LI() for _ in range(n))]
sushi.sort()
sushi.reverse()

# sushiを美味しいものからk個とる
umami_max = []  # 寿司の種類数が減らないように最低一つは別で確保
umami_sub = []  # 種類数を増やす時に捨てる候補(priority_queueで管理)
used = set()
for i in range(k):
    d, t = sushi[i]
    if t not in used:
        umami_max.append(d)
        used.add(t)
    else:
        heapq.heappush(umami_sub, d)

umami_1 = sum(umami_max)
umami_2 = sum(umami_sub)
sushi_v = len(used)
ans = umami_1 + umami_2 + sushi_v ** 2

# 捨てる寿司の候補がない場合，満足度はこれ以上大きくならない
if len(umami_sub) == 0:
    print(ans)
    exit()

for i in range(k, n):
    d, t = sushi[i]
    if t in used:
        continue
    used.add(t)
    umami_1 += d
    umami_2 -= heapq.heappop(umami_sub)  # 一番美味しくない寿司を消す
    sushi_v += 1
    ans = max(ans, umami_1 + umami_2 + sushi_v ** 2)

print(ans)

Submission status: Runtime Error

Solution code:
# coding:utf-8

import sys
import heapq
from collections import defaultdict, deque

INF = float('inf')
MOD = 10 ** 9 + 7

def LI(): return [int(x) for x in sys.stdin.readline().split()]
def LS(): return sys.stdin.readline().split()
def II(): return int(sys.stdin.readline())
def SI(): return input()


n, k = LI()
sushi = [(d, t) for t, d in (LI() for _ in range(n))]
sushi.sort()
sushi.reverse()

umami_max = []
umami_sub = []
used = set()
for i in range(k):
    d, t = sushi[i]
    if t not in used:
        umami_max.append(d)
        used.add(t)
    else:
        heapq.heappush(umami_sub, d)

umami_1 = sum(umami_max)
umami_2 = sum(umami_sub)
sushi_v = len(used)
ans = umami_1 + umami_2 + sushi_v ** 2
for i in range(k, n):
    d, t = sushi[i]
    if t in used:
        continue
    used.add(t)
    umami_1 += d
    umami_2 -= heapq.heappop(umami_sub)
    sushi_v += 1
    ans = max(ans, umami_1 + umami_2 + sushi_v ** 2)

print(ans)