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:
# -*- coding: utf-8 -*-

def popMaxOishisa(sushis):
    t, d = max(sushis["max"].items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


def popMinOishisa_nonUniqueNeta(sel_sushis):
    t, d = min(sel_sushis["min"].items(), key=lambda i:i[1])
    sel_sushis[t].pop()
    if len(sel_sushis[t]) == 1:
        sel_sushis["min"].pop(t)
    else:
        sel_sushis["min"][t] = sel_sushis[t][-1]
    return t, d


def popMaxOishisa_uniqueNeta(sushis, netas):
    max_dict = {k:v for k,v in sushis["max"].items() if not k in netas}
    t, d = max(max_dict.items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


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

# Store data
sushis = {}
sel_sushis = {}
for i in range(n):
    t, d = map(int, input().split())
    if not t in sushis:
        sushis[t] = [d]
    else:
        sushis[t].append(d)

# Formatting
for type in sushis:
    sushis[type].sort()
sushis["max"] = {t:sushis[t][-1] for t in sushis}

tMax = len(sushis) - 1

# Select sushi set as maximum oishisa point
sum_oishisa = 0
for i in range(k):
    t, d = popMaxOishisa(sushis)
    if not t in sel_sushis:
        sel_sushis[t] = [d]
    else:
        sel_sushis[t].append(d)
    sum_oishisa += d
sel_sushis["min"] = {t:sel_sushis[t][-1] for t in sel_sushis if len(sel_sushis[t]) > 1}

tMin = len(sel_sushis) - 1
max_score = tMin**2 + sum_oishisa

for nNeta in range(tMin+1, tMax + 1, 1):

    # Find sushi having minimum oishisa and being not unique for neta
    t1, d1 = popMinOishisa_nonUniqueNeta(sel_sushis)
    sum_oishisa -= d1

    # Find sushi having maximum oishisa and unique for neta
    t2, d2 = popMaxOishisa_uniqueNeta(sushis, list(sel_sushis.keys()))
    sum_oishisa += d2

    if not t1 in sushis:
        sushis[t1] = [d1]
    else:
        sushis[t1].append(d1)

    if not t2 in sel_sushis:
        sel_sushis[t2] = [d2]
    else:
        sel_sushis[t2].append(d2)

    score = nNeta**2 + sum_oishisa
    if score > max_score:
        max_score = score

print(max_score)

Submission status: Runtime Error

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

def popMaxOishisa(sushis):
    t, d = max(sushis["max"].items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


def popMinOishisa_nonUniqueNeta(sel_sushis):
    t, d = min(sel_sushis["min"].items(), key=lambda i:i[1])
    sel_sushis[t].pop()
    if len(sel_sushis[t]) == 1:
        sel_sushis["min"].pop(t)
    else:
        sel_sushis["min"][t] = sel_sushis[t][-1]
    return t, d


def popMaxOishisa_uniqueNeta(sushis, netas):
    max_dict = {k:v for k,v in sushis["max"].items() if not k in netas}
    t, d = max(sushis["max"].items(), key=lambda i:i[1])
    if len(sushis[t]) == 1:
        sushis.pop(t)
        sushis["max"].pop(t)
    else:
        sushis[t].pop()
        sushis["max"][t] = sushis[t][-1]
    return t, d


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

# Store data
sushis = {}
sel_sushis = {}
for i in range(n):
    t, d = map(int, input().split())
    if not t in sushis:
        sushis[t] = [d]
    else:
        sushis[t].append(d)

# Formatting
for type in sushis:
    sushis[type].sort()
sushis["max"] = {t:sushis[t][-1] for t in sushis}

tMax = len(sushis) - 1

# Select sushi set as maximum oishisa point
sum_oishisa = 0
for i in range(k):
    t, d = popMaxOishisa(sushis)
    if not t in sel_sushis:
        sel_sushis[t] = [d]
    else:
        sel_sushis[t].append(d)
    sum_oishisa += d
sel_sushis["min"] = {t:sel_sushis[t][-1] for t in sel_sushis if len(sel_sushis[t]) > 1}

tMin = len(sel_sushis) - 1
max_score = tMin**2 + sum_oishisa

for nNeta in range(tMin+1, tMax + 1, 1):

    # Find sushi having minimum oishisa and being not unique for neta
    t1, d1 = popMinOishisa_nonUniqueNeta(sel_sushis)
    sum_oishisa -= d1

    # Find sushi having maximum oishisa and unique for neta
    t2, d2 = popMaxOishisa_uniqueNeta(sushis, list(sel_sushis.keys()))
    sum_oishisa += d2

    if not t1 in sushis:
        sushis[t1] = [d1]
    else:
        sushis[t1].append(d1)

    if not t2 in sel_sushis:
        sel_sushis[t2] = [d2]
    else:
        sel_sushis[t2].append(d2)

    score = nNeta**2 + sum_oishisa
    if score > max_score:
        max_score = score

print(max_score)

Submission status: Runtime Error

Solution code:
package main

import (
    "bufio"
    "os"
    "strings"
    "strconv"
    "fmt"
    "sort"
    "container/heap"
    "math"
)

var sc = bufio.NewScanner(os.Stdin)

func readLineToInt() []int {
    sc.Scan()
    line := strings.Split(sc.Text(), " ")
    var intLine []int
    for _, str := range line {
        n, _ := strconv.Atoi(str)
        intLine = append(intLine, n)
    }
    return intLine
}


/****************** Pair ********************/
type Pair struct {
    a int
    b int
}

type SortablePairs struct {
    arr []*Pair
    lessfunc []lessFunc
}

type lessFunc func(i, j *Pair) bool

func (sp SortablePairs) Len() int {
    return len(sp.arr)
}

func (sp SortablePairs) Swap(i, j int) {
    sp.arr[i], sp.arr[j] = sp.arr[j], sp.arr[i]
}

func (sp SortablePairs) Less(i, j int) bool {
    p, q := sp.arr[i], sp.arr[j]
    for k := 0; k < len(sp.lessfunc); k++ {
        less := sp.lessfunc[k]
        if less(p, q) {
            return true
        } else if less(q, p) {
            return false
        }
    }
    return false
}

/***************** HeapQue *****************/

type PriorityQueue struct {
    arr []*Pair
    size int
}
/** sort **/
func (pq PriorityQueue) Len() int {
    return len(pq.arr)
}

func (pq PriorityQueue) Swap(i, j int) {
    pq.arr[i], pq.arr[j] = pq.arr[j], pq.arr[i]
}

func (pq PriorityQueue) Less(i, j int) bool{
    return pq.arr[i].a < pq.arr[j].a
}

/**** pq ****/
func (pq *PriorityQueue) Push(x interface{}) {
    item := x.(*Pair)
    pq.arr = append(pq.arr, item)
    pq.size++
}

func (pq *PriorityQueue) Pop() interface{} {
    n := pq.size
    item := pq.arr[n - 1]
    pq.arr = pq.arr[:n - 1]
    pq.size--
    return item
}

func max(a, b int) int {
    return int(math.Max(float64(a), float64(b)))
}

func main() {
    nk := readLineToInt()
    n, k := nk[0], nk[1]
    var dic map[int][]int = map[int][]int{}
    var keys []int
    for i := 0; i < n; i++ {
        td := readLineToInt()
        t, d := td[0], td[1]
        _, ok := dic[t]
        if !ok {
            keys = append(keys, t)
        }
        dic[t] = append(dic[t], d)
    }

    var sp SortablePairs
    less := func(p, q *Pair) bool {
        return p.a >= q.a
    }
    sp.lessfunc = append(sp.lessfunc, less)
    for _, key := range keys {
        sort.Ints(dic[key])
        length := len(dic[key])
        sp.arr = append(sp.arr, &Pair{dic[key][length - 1], key})
        dic[key] = dic[key][:length - 1]
    }
    sort.Sort(sp)
    var que PriorityQueue
    heap.Init(&que)
    queScore := 0
    defaultScore := 0
    cnt := 0
    ans := 0
    for _, p := range sp.arr {
        defaultScore += p.a + cnt * 2 + 1
        cnt++
        if cnt >= k {
            ans = max(ans, defaultScore)
            break
        }
        for _, add := range dic[p.b] {
            heap.Push(&que, &Pair{add, p.b})
            queScore += add
        }
        for que.size > k - cnt {
            item := heap.Pop(&que).(*Pair)
            queScore -= item.a
        }
        ans = max(ans, defaultScore + queScore)
    }
    fmt.Println(ans)

}

Submission status: Runtime Error

Solution code:
from math import pow
from collections import deque, Counter
# from loguru import logger


def cal_score(N, K, records):
    sorted_recs = sorted(records, key=lambda x: -x[1])
    target_queue = deque(sorted_recs[0:K]) # only allow pop() action, so use dequeue
    # logger.debug(target_queue)
    compare_queue = deque(sorted_recs[K:])
    count = Counter()
    for neta_point in target_queue:
        count[neta_point[0]] += 1
    # logger.debug(count)

    cardinality = len(count.keys())
    sum_target_points = sum([n_p[1] for n_p in target_queue])
    score = pow(cardinality, 2) + sum_target_points
    score = int(score)
    new_scores = deque()
    new_scores.append(score)
    # start iteration and comparing gain and lost
    while compare_queue:
        comp = compare_queue.popleft()
        if comp[0] in count:
            continue
        while target_queue:
            old = target_queue.pop()
            if count[old[0]] < 2:
                # swap this neta will not increase cardinality but lower neta point
                continue
            else: # check if score will go up
                inc = 2 * cardinality + 1
                dec = old[1] - comp[1]
                delta = inc - dec
                # logger.debug(f"delta is {delta}")
                # if delta >= 0:
                #     logger.debug("SOME")
                #     logger.debug(f"{old} out")
                #     logger.debug(f"{comp} in")
                #     cardinality += 1
                #     score += delta
                #     count[old[0]] -= 1
                #     count[comp[0]] +=1
                # else: #increase cardinality will bring more lost than gain
                #     # exit here
                #     return score
                new_score = new_scores[-1] + delta
                new_scores.append(new_score)

                cardinality += 1
                count[old[0]] -= 1
                count[comp[0]] +=1
    # both queues are drained
    return max(new_scores)
N, K = map(int, input().split())
records = [map(int, input().split()) for i in range(N)]
score = cal_score(N, K, records)
print(score)

Submission status: Runtime Error

Solution code:
import sys
stdin = sys.stdin
 
sys.setrecursionlimit(10**5) 
 
def li(): return map(int, stdin.readline().split())
def li_(): return map(lambda x: int(x)-1, stdin.readline().split())
def lf(): return map(float, stdin.readline().split())
def ls(): return stdin.readline().split()
def ns(): return stdin.readline().rstrip()
def lc(): return list(ns())
def ni(): return int(stdin.readline())
def nf(): return float(stdin.readline())

n,k = li()
td = [tuple(li()) for _ in range(n)]

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

kindset = set()
doublestack = []

ans = 0

for ti, di in td[:k]:
    if ti in kindset:
        doublestack.append((ti,di))
    
    else:
        kindset.add(ti)
        
    ans += di

ans += len(kindset) * len(kindset)
score = ans

# n-k個の要素に対して重複を排して交換
for ti, di in td[k:]:
    if ti in kindset:
        continue
    
    #スコア再計算
    old_ti, old_di = doublestack.pop()    
    score = score - old_di + di + 2*len(kindset) + 1
    
    # 更新
    ans = max(ans, score)    
    kindset.add(ti)
    
    # すべてユニークになったら抜ける
    if not len(doublestack):
        break
    
print(ans)

Submission status: Runtime Error

Solution code:
import sys
stdin = sys.stdin
 
sys.setrecursionlimit(10**5) 
 
def li(): return map(int, stdin.readline().split())
def li_(): return map(lambda x: int(x)-1, stdin.readline().split())
def lf(): return map(float, stdin.readline().split())
def ls(): return stdin.readline().split()
def ns(): return stdin.readline().rstrip()
def lc(): return list(ns())
def ni(): return int(stdin.readline())
def nf(): return float(stdin.readline())

n,k = li()
td = [tuple(li()) for _ in range(n)]

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

kindset = set()
doublestack = []

ans = 0

for ti, di in td[:k]:
    if ti in kindset:
        doublestack.append((ti,di))
    
    else:
        kindset.add(ti)
        
    ans += di

ans += len(kindset) * len(kindset)
score = ans

# n-k個の要素に対して重複を排して交換
for ti, di in td[k:]:
    if ti in kindset:
        continue
    
    #スコア再計算
    old_ti, old_di = doublestack.pop()    
    score = score - old_di + di + 2*len(kindset) + 1
    
    # 更新
    ans = max(ans, score)    
    kindset.add(ti)
    
    # すべてユニークになったら抜ける
    if len(doublestack):
        break
    
print(ans)

Submission status: Runtime Error

Solution code:
import sys
stdin = sys.stdin
 
sys.setrecursionlimit(10**5) 
 
def li(): return map(int, stdin.readline().split())
def li_(): return map(lambda x: int(x)-1, stdin.readline().split())
def lf(): return map(float, stdin.readline().split())
def ls(): return stdin.readline().split()
def ns(): return stdin.readline().rstrip()
def lc(): return list(ns())
def ni(): return int(stdin.readline())
def nf(): return float(stdin.readline())

n,k = li()
td = [tuple(li()) for _ in range(n)]

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

kindset = set()
doublestack = []

ans = 0

for ti, di in td[:k]:
    if ti in kindset:
        doublestack.append((ti,di))
    
    else:
        kindset.add(ti)
        
    ans += di

ans += len(kindset) * len(kindset)
score = ans

# n-k個の要素に対して重複を排して交換
for ti, di in td[k:]:
    if ti in kindset:
        continue
    
    old_ti, old_di = doublestack.pop()
    
    score = score - old_di + di + 2*len(kindset) + 1
    
    ans = max(ans, score)
    
    kindset.add(ti)
    
print(ans)

Submission status: Runtime Error

Solution code:
from collections import defaultdict
from heapq import heappush, heappop


n, k = map(int, input().split())
dic = defaultdict(list)
for _ in range(n):
    t, d = map(int, input().split())
    dic[t].append(d)

maxs = []
for key in dic:
    dic[key].sort()
    maxs.append((dic[key].pop(), key))

maxs.sort(reverse=True)
queSize = 0
que = []
queScore = 0
defaultScore = 0
cnt = 0
ans = 0
for v, key in maxs:
    defaultScore += v + cnt * 2 + 1
    cnt += 1
    for add in dic[key]:
        heappush(que, add)
        queSize += 1
        queScore += add
    while queSize > k - cnt:
        queScore -= heappop(que)
        queSize -= 1
    ans = max(ans, defaultScore + queScore)

print(ans)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat)+1)):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat))+1)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Submission status: Runtime Error

Solution code:
from collections import defaultdict as dd
import heapq

n,k=map(int,input().split())
l=[list(map(int,input().split())) for _ in range(n)]
l.sort(key=lambda x:-x[1])

sushi_types = set()
sushies = []
v=0

for i in range(k):
    t,d=l[i]
    v+=d
    if t in sushi_types:
        heapq.heappush(sushies,d)
    else:
        sushi_types.add(t)

remaining_sushi=dd(list)
for i in range(k,n):
    t,d=l[i]
    if not t in sushi_types:
        heapq.heappush(remaining_sushi[t],-d)

l=[0]*(k+1)
current_types=len(sushi_types)
l[current_types]=v+current_types**2
ans=l[current_types]

for i in range(len(sushi_types)+1,k+1):
    t,add_target=min(remaining_sushi.items(),key=lambda x:x[1][0])
    sushi_types.add(t)
    remove_target=heapq.heappop(sushies)
    l[i]=l[i-1]-(i-1)**2-add_target[0]-remove_target+i**2
    ans=max(ans,l[i])
    del remaining_sushi[t]

print(ans)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat)))]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
        ans = ans-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
    else:
        break
print(ans)

Submission status: Runtime Error

Solution code:
N,K= map(int,input().split())
n=0
lis=[]
points=0
kind=set()
l1=[]
l2=[]
for _ in range(N):
    t,d = map(int,input().split())
    lis.append((t,d))
lis.sort(key=lambda x:x[1],reverse=True)
for i in lis:
    if K>n:
        if i[0] not in kind:
            l1.append(i[1])
            kind.add(i[0])
        else:
            l2.append(i[1])
        n+=1
        if K==n:
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i[1])
            kind.add(i[0])
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
            if len(l2)==0:
                break

print(points)

Submission status: Runtime Error

Solution code:
N,K= map(int,input().split())
n=0
lis=[]
points=0
kind=set()
kindall=set()
l1=[]
l2=[]
for _ in range(N):
    t,d = map(int,input().split())
    kindall.add(t)
    lis.append((t,d))
lis.sort(key=lambda x:x[1],reverse=True)
for i in lis:
    if K>n:
        if i[0] not in kind:
            l1.append(i[1])
            kind.add(i[0])
        else:
            l2.append(i[1])
        n+=1
        if K==n:
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i[1])
            kind.add(i[0])
            tl=sum(l1) +sum(l2) +len(kind)**2
            points=max(tl,points)
            
        if len(l2)==0:
            break

print(points)

Submission status: Runtime Error

Solution code:
def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

N,K=map(int, input().split())
L=[list(map(int, input().split())) for _ in range(N)]

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

L_kind=set([])
L_sorted_kind=set([])
bucket={}
L_sorted_=[]
tmp_set=set([])
for i,x in enumerate(L):
  L_kind|=set([x[0]])
  if i<K:
    L_sorted_kind|=set([x[0]])
    if x[0] in bucket:
      bucket[x[0]]+=1
    else:
      bucket[x[0]]=1
  else:
    if not (x[0] in L_sorted_kind) and not (x[0] in tmp_set):
      tmp_set|=set([x[0]])
      L_sorted_.append(x)

max_score=score(L_sorted[:K])

i=1
while len(L_kind & L_sorted_kind)!=len(L_kind) and i<=K:
  if bucket[L_sorted[:K][-i][0]]>1:
    bucket[L_sorted[:K][-i][0]]-=1
    
    L_sorted[K-i]=L_sorted_[0]
    L_sorted_.pop(0)
    
    L_sorted_kind=set([i[0] for i in L_sorted[:K]])

    tmp_score=score(L_sorted[:K])
    if tmp_score>max_score:
      max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Runtime Error

Solution code:
def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

N,K=map(int, input().split())
L=[list(map(int, input().split())) for _ in range(N)]

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

L_kind=[]
L_sorted_kind=[]
bucket={}
L_sorted_=[]
tmp_set=[]
for i,x in enumerate(L):
  L_kind.append(x[0])
  if i<K:
    L_sorted_kind.append(x[0])
    if x[0] in bucket:
      bucket[x[0]]+=1
    else:
      bucket[x[0]]=1
  else:
    if not (x[0] in set(L_sorted_kind)) and not (x[0] in set(tmp_set)):
      tmp_set.append(x[0])
      L_sorted_.append(x)
L_kind=set(L_kind)
L_sorted_kind=set(L_sorted_kind)

max_score=score(L_sorted[:K])

i=1
while len(L_kind & L_sorted_kind)!=len(L_kind) and i<=K:
  if bucket[L_sorted[:K][-i][0]]>1:
    bucket[L_sorted[:K][-i][0]]-=1
    
    L_sorted[K-i]=L_sorted_[0]
    L_sorted_.pop(0)
    
    L_sorted_kind=set([i[0] for i in L_sorted[:K]])

    tmp_score=score(L_sorted[:K])
    if tmp_score>max_score:
      max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
ans = 0
eat,noteat =[],[]
a = [0 for i in range(N)]
k = K
for i in range(N):
    if k == -1:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans += sum(eat)
ans += sum(noteat[:k])
ans += pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
        ans = ans-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
    else:
        break
print(ans)

Submission status: Runtime Error

Solution code:
N,K=map(int,input().split())
s=[]
S=[]
ss=[]
for i in range(N):
    t,d=map(int,input().split())
    s.append([d,t])
s.sort(reverse=True)
for i in range(N):
    S.append(s[i][1])
T=[]
Ta=[]
Tb=[]
for i in range(N):
    if i>=K:
        if S[i] in S[:i]:
            T.append(i)
    else:
        if S[i] in S[:i]:
            ss.append(i)
for i in range(N):
    if i not in T:
        Ta.append(s[i])
        Tb.append(S[i])
s=Ta
S=Tb
a=0
b=0
for j in range(K):
    a+=s[j][0]
    if s[j][1] not in S[:j]:
        b+=1
a+=b*b
for k in range(len(ss)):
    a1=0
    b1=0
    S=S[:ss[-k-1]]+S[ss[-k-1]+1:]
    s=s[:ss[-k-1]]+s[ss[-k-1]+1:]
    for j in range(K):
        a1+=s[j][0]
        if s[j][1] not in S[:j]:
            b1+=1
    if a1+b1*b1>a:
        a=a1+b1*b1
print(a)

Submission status: Runtime Error

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

def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2


L=[list(map(int, input().split())) for _ in range(N)]
L_sorted=sorted(L, key=lambda x:x[1], reverse=True)

L_kind=set([i[0] for i in L])
L_sorted_kind=set([i[0] for i in L_sorted[:K]])


bucket={k:len([k for i in L if i[0]==k]) for k in L_sorted_kind}

L_sorted_=L_sorted[K:]
L_sorted_C=[i for i in L_sorted_ if len(set([i[0]]) & set(L_sorted_kind))==0]

max_score=score(L_sorted[:K])

i=1
while len(L_kind & L_sorted_kind)!=len(L_kind) and i<=K:
  if bucket[L_sorted[:K][-i][0]]>1:
    bucket[L_sorted[:K][-i][0]]-=1
    
    L_sorted[K-i]=L_sorted_C[0]
    L_sorted_C.pop(0)
    
    L_sorted_kind=set([i[0] for i in L_sorted[:K]])

  tmp_score=score(L_sorted[:K])
  if tmp_score>max_score:
    max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Runtime Error

Solution code:
N,K=map(int,input().split())
s=[]
S=[]
ss=[]
for i in range(N):
    t,d=map(int,input().split())
    s.append([d,t])
s.sort(reverse=True)
for i in range(N):
    S.append(s[i][1])
T=s
Ta=S
for i in range(N):
    if i>=K:
        if S[i] in S[:i]:
            T=s[:i]+s[i+1:]
            Ta=S[:i]+s[i+1:]
    else:
        if S[i] in S[:i]:
            ss.append(i)
s=T
S=Ta
a=0
b=0
for j in range(K):
    a+=s[j][0]
    if s[j][1] not in S[:j]:
        b+=1
a+=b*b
for k in range(len(ss)):
    a1=0
    b1=0
    S=S[:ss[-k-1]]+S[ss[-k-1]+1:]
    s=s[:ss[-k-1]]+s[ss[-k-1]+1:]
    for j in range(K):
        a1+=s[j][0]
        if s[j][1] not in S[:j]:
            b1+=1
    if a1+b1*b1>a:
        a=a1+b1*b1
print(a)

Submission status: Runtime Error

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

def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

L=[list(map(int, input().split())) for _ in range(N)]

L_sorted=sorted(L, key=lambda x:x[1], reverse=True)
L_kind=set([i[0] for i in L])
L_kind_=set([i[0] for i in L_sorted[:K]])

tmp={k:len([k for i in L if i[0]==k]) for k in L_kind}

max_score=score(L_sorted[:K])

L_sorted_kind_=[L_sorted[K:][l] for l in range(len(L)-K) if len(set([L_sorted[K:][l][0]]) & set(L_kind_))==0]

i=1
while len(L_kind & L_kind_)!=len(L_kind) and i<=K:
  if tmp[L_sorted[:K][-i][0]]>1:
    tmp[L_sorted[:K][-i][0]]-=1
    L_sorted[K-i]=L_sorted_kind_[0]
    L_sorted_kind_.pop(0)

  tmp_score=score(L_sorted[:K])
  if tmp_score>max_score:
    max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Runtime Error

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 21:13:02 2019

@author: Yamazaki Kenichi
"""

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

A.sort(key=lambda x: x[1])
kinds = 0
t = 0
for i in range(N):
    if t != A[i][0]:
        kinds += 1
        t = A[i][0]
c = K - kinds if K - kinds >=0 else 0
d = K-c
temp_t = pow(K-c,2)
i = N-1
counta =0
temp_d=0
a = [0 for i in range(N+1)]
b,b2 =[],[]
while counta < K:
    if c-a[A[i][0]] >= 0:
        temp_d += A[i][1]
        c = c-1 if a[A[i][0]] >= 1 else c
        a[A[i][0]] += 1
        counta+=1
        b2.append(A[i][1])
    else:
        b.append(A[i][1])
    i -=1
ans = [0 for i in range(len(b)+1)]
ans[0] = temp_t+temp_d
#print(ans)
for i in range(len(b)):
    ans[i+1] = ans[i] +b[i]-b2[-1-i]-(pow(d,2)-pow(d-i-1,2))
print(max(ans))

Submission status: Runtime Error

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

def score(l):
  return sum([i[1] for i in l])+len(set([i[0] for i in l]))**2

L=[list(map(int, input().split())) for _ in range(N)]

L_sorted=sorted(L, key=lambda x:x[1], reverse=True)
L_kind=set([i[0] for i in L])
L_kind_=set([i[0] for i in L_sorted[:K]])

tmp={k:0 for k in L_kind}

for k in L_kind:
  tmp[k]=len([k for i in L if i[0]==k])

max_score=score(L_sorted[:K])

i=1
while len(L_kind & L_kind_)!=len(L_kind) and i<=K:
  if tmp[L_sorted[:K][-i][0]]>1:
    tmp[L_sorted[:K][-i][0]]-=1
    for l in range(len(L)-K):
      ttt=L_sorted[K:][l]
      if len(set([ttt[0]]) & set(L_kind_))==0:
        L_sorted[K-i]=ttt
        L_sorted.pop(K+l)
        break

  tmp_score=score(L_sorted[:K])
  if tmp_score>max_score:
    max_score=tmp_score
  i+=1    
  
print(max_score)

Submission status: Runtime Error

Solution code:
import copy
import collections


def score(list):
    tl, dl = zip(*list)
    t_s = len(set(tl)) ** 2
    d_s = sum(dl)
    return t_s + d_s


def r(n, k, sl):
    
    tl, dl = zip(*sl)
    ztd = collections.Counter(tl)
    for _ in range(len(sl) - k):
        d_idx = dl.index(min(dl))
        tlv = tl[d_idx]
        if ztd[tlv] >= 2:
            del sl[d_idx]
        else:
            break

    for _ in range(len(sl) - k):

        ckl = {}
        for v in sl:
            if v[0] in ckl:
                if ckl[v[0]][1] > v[1]:
                    ckl[v[0]] = v
            else:
                ckl[v[0]] = v

        s_c = 0
        dtl = []
        for delckl in ckl.values():
            c_sl = copy.copy(sl)
            del c_sl[c_sl.index(delckl)]
            soe = score(c_sl)
            if s_c < soe:
                s_c = soe
                dtl = delckl

        del sl[sl.index(dtl)]

    return score(sl)


def main():
    n, k = list(map(int, input().split()))
    sl = [list(map(int, input().split())) for i in range(n)]
    print(r(n, k, sl))


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
N, K = (int(i) for i in input().split())
L = []
for i in range(N):
    t,d = (int(i) for i in input().split())
    L.append((t,d))

L = sorted(L, key = lambda x:(x[0], -x[1]))
for i in range(N):
    if i == 0:
        ra = 0
        L[i] = (L[i][0], L[i][1], ra)
    else:
        if L[i][0] == L[i-1][0]:
            ra += 1
        else:
            ra = 0
        L[i] = (L[i][0], L[i][1], ra)

L = sorted(L, key = lambda x:(-x[1], x[0]))

if N != K:
    kouho = L[:K]
    amari = L[K:]
else:
    kouho = L
    
nowval = len(set([sushi[0] for sushi in kouho]))
nowder = sum([sushi[1] for sushi in kouho])
point = nowval ** 2 + nowder

if N == K:
    print(point)
else:
    weaks = []
    for sushi in kouho:
        if sushi[2] > 0:
            weaks.append(sushi)
    weaks.reverse()
    kaburan = []
    for sushi in amari:
        if sushi[2] == 0:
            kaburan.append(sushi)
    
    ma = max(len(weaks), len(kaburan))
    newpoint = point
    for i in range(ma):
        newpoint = newpoint - weaks[i][1] + kaburan[i][1]
        newpoint = newpoint - nowval ** 2 + (nowval + 1) ** 2
        nowval += 1
        point = max(point, newpoint)
    print(point)