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: Time Limit Exceeded

Solution code:
from bisect import bisect_left, bisect_right
from collections import Counter, defaultdict, deque
from copy import deepcopy
from functools import lru_cache, reduce
from math import ceil, floor

import heapq
import itertools
import operator


def get_int():
    return int(input())


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


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


def get_char_list():
    return list(input().strip())


def satisfy(sushis):
    t_count = len({t for i, t, d in sushis})
    d_sum = sum(d for i, t, d in sushis)
    return t_count * t_count + d_sum


def main():
    n, k = get_int_list()
    sushis = []
    for i in range(n):
        t, d = get_int_list()
        sushis.append((i, t, d))

    ss = deepcopy(sushis)
    s_dict = defaultdict(list)
    for t, s in itertools.groupby(ss, key=lambda x: x[1]):
        s_dict[t].extend(list(s))
    max_d_list = []
    remains = []
    for t in s_dict.keys():
        s_dict[t].sort(key=lambda x: x[2])
        max_d_list.append(s_dict[t].pop(-1))
        remains.extend(s_dict[t])
    max_d_list.sort(key=lambda x: x[2], reverse=True)
    remains.sort(key=lambda x: x[2], reverse=True)
    bests = []
    t_count = len(list(s_dict.keys()))
    if k > t_count:
        bests = sorted(max_d_list[:] + remains[:k-t_count], key=lambda x: x[2])
        remains = reversed(remains[k-t_count:])
    else:
        bests = reversed(max_d_list[:k])
        remains = sorted(max_d_list[k:] + remains, key=lambda x: x[2])
    bests = list(bests)
    remains = list(remains)
    while remains:
        another = bests[1:] + remains[-1:]
        if satisfy(bests) > satisfy(another):
            break
        bests = another
        remains.pop(-1)
    print(satisfy(bests))


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
# -*- coding: utf-8 -*-
import itertools
import collections

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

for i in range(N):
  shurui, oisisa=map(int, input().split())
  shuruiList.append(shurui)
  oisisaKisoList.append(oisisa)

l = range(0, N, 1)
p_list = list(itertools.combinations(l, K))

max=0
for i in range(len(p_list)):
  shuruiListTmp=[]
  oisisaKisoTmp=[]
  maxTmp=0
  for j in p_list[i]:
    shuruiListTmp.append(shuruiList[j])
    maxTmp+=oisisaKisoList[j]
    
  shuruipoint=len(collections.Counter(shuruiListTmp))
  maxTmp+=shuruipoint*shuruipoint
  
  if maxTmp>max:
    max=maxTmp
    
print(max)

Submission status: Time Limit Exceeded

Solution code:
from sys import stdin
NK = [int(i) for i in stdin.readline().rstrip().split(" ")]
T = []
D = []
for i in range(NK[0]):
    tmp = [int(i) for i in stdin.readline().rstrip().split(" ")]
    T.append(tmp[0])
    D.append(tmp[1])

last = NK[0] - NK [1]
def remover(dishesT,dishesD,currentLast):
    variety = len(set(dishesT))
    if currentLast == 0:
        return sum(dishesD) + variety ** 2
    else:
        surplus = [-i - (2*variety)+1 for i in dishesD]
        for j in range(len(dishesT)):
            if dishesT[j] in [x for i,x in enumerate(dishesT) if i!=j]:
                surplus[j] += (2*variety -1)
        indexer = surplus.index(max(surplus))
        return remover([x for i,x in enumerate(dishesT) if i!=indexer],
        [x for i,x in enumerate(dishesD) if i!=indexer],
        currentLast -1)

print(remover(T,D,last))

Submission status: Time Limit Exceeded

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 20:52:35 2019

@author: Owner
"""

import collections
import scipy.misc
import sys
import numpy as np
import math
from operator import itemgetter
import itertools
import copy

#素因数を並べる
def prime_decomposition(n):
  i = 2
  table = []
  while i * i <= n:
    while n % i == 0:
      n /= i
      table.append(int(i))
    i += 1
  if n > 1:
    table.append(int(n))
  return table   
# 桁数を吐く
def digit(i):
    if i > 0:
        return digit(i//10) + [i%10]
    else:
        return []
    
def getNearestValueIndex(list, num):
    """
    概要: リストからある値に最も近い値のインデックスを取得する関数
    @param list: データ配列
    @param num: 対象値
    @return 対象値に最も近い値
    """

    # リスト要素と対象値の差分を計算し最小値のインデックスを取得
    idx = np.abs(np.asarray(list) - num).argmin()
    return idx

"""
N, X = map(int, input().split())

x = [0]*N
x[:] = map(int, input().split())

P = [0]*M
Y = [0]*M
for m in range(M):
    P[m], Y[m] = map(int, input().split())
    
all(nstr.count(c) for c in '753')
"""

N, K = map(int, input().split())

t = [0]*N
d = [0]*N
for n in range(N):
    t[n], d[n] = map(int, input().split())
a = [i for i in range(N)]
L = list(itertools.combinations(a,K))
result=0
for l in range(len(L)):
    tmp_t=[]
    oisisa = 0
    tmp_result = 0
    li = L[l]
    for idx in li:
        tmp_t.append(t[idx])
        oisisa += d[idx]
    
    X = collections.Counter(tmp_t)
    num_key = len(X.keys())
    
    tmp_result = num_key**2 + oisisa
    if tmp_result > result:
        result = tmp_result
print(result)

Submission status: Time Limit Exceeded

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 20:52:35 2019

@author: Owner
"""

import collections
import scipy.misc
import sys
import numpy as np
import math
from operator import itemgetter
import itertools
import copy

#素因数を並べる
def prime_decomposition(n):
  i = 2
  table = []
  while i * i <= n:
    while n % i == 0:
      n /= i
      table.append(int(i))
    i += 1
  if n > 1:
    table.append(int(n))
  return table   
# 桁数を吐く
def digit(i):
    if i > 0:
        return digit(i//10) + [i%10]
    else:
        return []
    
def getNearestValueIndex(list, num):
    """
    概要: リストからある値に最も近い値のインデックスを取得する関数
    @param list: データ配列
    @param num: 対象値
    @return 対象値に最も近い値
    """

    # リスト要素と対象値の差分を計算し最小値のインデックスを取得
    idx = np.abs(np.asarray(list) - num).argmin()
    return idx

"""
N, X = map(int, input().split())

x = [0]*N
x[:] = map(int, input().split())

P = [0]*M
Y = [0]*M
for m in range(M):
    P[m], Y[m] = map(int, input().split())
    
all(nstr.count(c) for c in '753')
"""

N, K = map(int, input().split())

t = [0]*N
d = [0]*N
for n in range(N):
    t[n], d[n] = map(int, input().split())
a = [i for i in range(N)]
b = [i for i in range(K)]
L = list(itertools.combinations(a,K))
result=0
for l in range(len(L)):
    tmp_t=[]
    oisisa = 0
    tmp_result = 0
    li = L[l]
    for idx in li:
        tmp_t.append(t[idx])
        oisisa += d[idx]
    
    X = collections.Counter(tmp_t)
    num_key = len(X.keys())
    
    tmp_result = num_key**2 + oisisa
    if tmp_result > result:
        result = tmp_result
print(result)

Submission status: Time Limit Exceeded

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)
            kind.add(i[0])
        else:
            l2.append(i)
        n+=1
        tl=0
        for k in l1:
            tl+=k[1]
        for k in l2:
            tl+=k[1]
        tl+=len(kind)**2
        points=max(tl,points)
    else:
        if i[0] in kind:
            continue
        else:
            l2.pop(-1)
            l1.append(i)
            kind.add(i[0])
            tl=0
            for k in l1:
                tl+=k[1]
            for k in l2:
                tl+=k[1]
            tl+=len(kind)**2
            points=max(tl,points)
            
    if len(kind)==K:
        break

print(points)
            #tmp=len(kind)

#             if l2[-1][1]<i[1]+(tmp+1)**2-tmp**2:
#                 l2.pop(-1)
#                 l1.append(i)
#                 kind.add(i[0])
# print(lis,l1,l2,kind)

Submission status: Time Limit Exceeded

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

now = data[:K]


def count(x):
    p = len(set([a[0] for a in x]))
    point = p ** 2
    p = sum([a[1] for a in x])
    point += p
    return point

p_max = count(now)

for i in list(itertools.combinations(data,K)):
    if p_max < count(i):
        p_max = count(i)


print(p_max)

Submission status: Time Limit Exceeded

Solution code:
import copy
import itertools

N, K = (int(x) for x in input().split())
t=[]
d=[]
for i in range(N):
    tmpT, tmpD = (int(x) for x in input().split())
    t.append(tmpT)
    d.append(tmpD)

combList = list(itertools.combinations(range(N), K))
manzokudo=[0]*len(combList)
for combID in range(len(combList)):
    comb = combList[combID]
    x=[]
    for i in range(K):
        manzokudo[combID] += d[comb[i]]
        if t[comb[i]] in x:
            pass
        else:
            x.append(t[comb[i]])
    manzokudo[combID] += len(x)*len(x)

print(max(manzokudo))

Submission status: Time Limit Exceeded

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  A.append(list(map(int,input().split())))

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
l = 0
k = False
D = []
while l < K:
  while True:
    j += 1
    if j >= len(A):
      j = 0
      k = True
    if k or not A[j][0] in D:
      ans2 += A[j][1]
      D.append(A[j][0])
      del A[j]
      j -= 1
      l += 1
      break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

Submission status: Time Limit Exceeded

Solution code:
from decimal import *
import itertools


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):
    s = 0
    for v in itertools.combinations(sl, k):
        cv = score(v)
        if cv > s:
            s = cv
    return s


def main():
    n, k = list(map(int, input().split()))
    sl = []
    for i in range(n):
        t, d = list(input().split())
        sl.append([int(t), Decimal(d)])

    print(r(n, k, sl))


if __name__ == '__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
import itertools
sushis = []
def satis(ixs):
    b = 0
    types = []
    for ix in ixs:
        b += sushis[ix][1]
        t = sushis[ix][0]
        if t not in types:
            types.append(t)
    return b + (len(types) ** 2)

n, k = map(int, input().split(' '))
for i in range(n):
    t, d = map(int, input().split(' '))
    sushis.append((t, d))
print(max(map(satis, (itertools.combinations(range(n), k)))))

Submission status: Time Limit Exceeded

Solution code:
import sys
input = sys.stdin.readline

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

typenum1 = 0
typehistory = []

for _ in range(N):
    t, d = map(int, input().split())
    dt.append([d, t])
    if t not in typehistory:
        typenum1 += 1
        typehistory.append(t)

dt.sort()
dt.reverse()
typenum2 = 0
array = [0] * (N+1)
result = 0

for i in range(K):
    result += dt[i][0]
    array[dt[i][1]] += 1
    if array[dt[i][1]] == 1:
        typenum2 += 1

result += typenum2 ** 2
num = result

while typenum2 < typenum1 and typenum2 < K:
    for i in range(K-1, -1, -1):
        if array[dt[i][1]] > 1:
            num -= dt[i][0]
            array[dt[i][1]] -= 1
            break

    for i in range(K, N):
        if array[dt[i][1]] == 0:
            num += dt[i][0]
            array[dt[i][1]] += 1
            break

    num += typenum2 * 2 + 1

    result = max(result, num)

    typenum2 += 1

print(result)

Submission status: Time Limit Exceeded

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

typenum1 = 0
typehistory = []

for _ in range(N):
    t, d = map(int, input().split())
    dt.append([d, t])
    if t not in typehistory:
        typenum1 += 1
        typehistory.append(t)

dt.sort()
dt.reverse()
typenum2 = 0
array = [0 for _ in range(N + 1)]
result = 0

for i in range(K):
    result += dt[i][0]
    array[dt[i][1]] += 1
    if array[dt[i][1]] == 1:
        typenum2 += 1

result += typenum2 ** 2
num = result

while typenum2 < typenum1 and typenum2 < K:
    for i in range(K-1, -1, -1):
        if array[dt[i][1]] > 1:
            num -= dt[i][0]
            array[dt[i][1]] -= 1
            break

    for i in range(K, N):
        if array[dt[i][1]] == 0:
            num += dt[i][0]
            array[dt[i][1]] += 1
            break

    num += typenum2 * 2 + 1

    result = max(result, num)

    typenum2 += 1

print(result)

Submission status: Time Limit Exceeded

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  IN = list(map(int,input().split()))
  A.append(IN)

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
k = False
D = []
for i in range(K):
  while True:
    j += 1
    if j >= len(A):
      j = 0
      k = True
    if k or not A[j][0] in D:
      ans2 += A[j][1]
      D.append(A[j][0])
      del A[j]
      j -= 1
      break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
neta = [[] for i in range(N+1)]
for i in range(N):
    t, d = map(int, input().split())
    neta[t].append(d)
for i in range(N):
    neta[i].sort(reverse = True)
neta = [i for i in neta if i != []]
neta.sort(reverse = True)
kinds = len(neta)
maxneta = 0
for i in range(len(neta)):
    maxneta = max(maxneta, len(neta[i]))
anslist = []
kinds_max = min(kinds, K)
x = K
for i in range(len(neta)):
    x -= len(neta[i])
    if x <= 0:
        kinds_min = i+1
        break
firstnetasum = [0]
for i in range(kinds):
    firstnetasum.append(firstnetasum[i] + neta[i][0])
for i in range(kinds_min, kinds_max + 1):
    ans = firstnetasum[i]
    remain = []
    for j in range(i):
        try:
            remain += neta[j][1:]
        except IndexError:
            pass
    remain.sort(reverse = True)
    ans += sum(remain[:K-i])
    ans += i**2
    anslist.append(ans)
print(max(anslist))

Submission status: Time Limit Exceeded

Solution code:
from itertools import permutations

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

ans = 0
for items in permutations(inputs, K):
    tmp = 0
    sets = set()
    for i in range(K):
        tmp += items[i][1]
        sets.add(items[i][0])
    tmp += len(sets) * len(sets)
    ans = max(ans, tmp)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
import math
import itertools

N,K = map(int, input().split())

t = []
d = []
neta = []
bass_score = 0

for _ in range(N):
    tmp1,tmp2 = map(int, input().split())
    t.append(tmp1)
    d.append(tmp2)


dd = d.copy()

for _ in range(K):
    index = dd.index(max(dd))
    value = max(dd)
    bass_score += value
    # print(value,index)
    neta.append(t[index])
    dd[index] = 0

bonus_score = len(set(neta))**2

# print("bass:",bass_score)
# print("bonus:",bonus_score)
# print("total:",bass_score + bonus_score)
test1 = bass_score + bonus_score
#----------------------------
dd = d.copy()
neta = []
bass_score = 0
count = 0
while True:
    index = dd.index(max(dd))
    if t[index] in neta:
        dd[index] = 0
        continue
    else:
        value = max(dd)
        bass_score += value
        # print(value,index)
        neta.append(t[index])
        dd[index] = 0
        count += 1
        if count == K:
            break

bonus_score = len(set(neta))**2

# print("bass:",bass_score)
# print("bonus:",bonus_score)
# print("total:",bass_score + bonus_score)
test2 = bass_score + bonus_score

print(max(test1,test2))

Submission status: Time Limit Exceeded

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  IN = list(map(int,input().split()))
  A.append(IN)

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
k = False
D = []
for i in range(K):
  while True:
    j += 1
    if j >= len(A):
      j = 0
      k = True
    if k:
      ans2 += A[j][1]
      if not A[j][0] in D:
        D.append(A[j][0])
      del A[j]
      j -= 1
      break
    if not A[j][0] in D:
      ans2 += A[j][1]
      if not A[j][0] in D:
        D.append(A[j][0])
      del A[j]
      j -= 1
      break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

Submission status: Time Limit Exceeded

Solution code:
N, K = map(int, input().split())
neta = [[] for i in range(N+1)]
for i in range(N):
    t, d = map(int, input().split())
    neta[t].append(d)
for i in range(N):
    neta[i].sort(reverse = True)
neta = [i for i in neta if i != []]
neta.sort(reverse = True)
kinds = len(neta)
maxneta = 0
for i in range(len(neta)):
    maxneta = max(maxneta, len(neta[i]))
anslist = []

kinds_max = min(kinds, K)
x = K
for i in range(len(neta)):
    x -= len(neta[i])
    if x <= 0:
        kinds_min = i+1
        break
firstnetasum = [0]
for i in range(kinds):
    firstnetasum.append(firstnetasum[i] + neta[i][0])
for i in range(kinds_min, kinds_max + 1):
    ans = firstnetasum[i]
    remain = []
    for j in range(i):
        try:
            remain += neta[j][1:]
        except IndexError:
            pass
    remain.sort(reverse = True)
    ans += sum(remain[:K-i])
    ans += i**2
    anslist.append(ans)
print(max(anslist))

Submission status: Time Limit Exceeded

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

def satisfaction(t, d, c):
    neta = set()
    kiso = 0
    for i in c:
        neta.add(t[i])
        kiso += d[i]
    x = len(neta)
    shurui = x * x

    return kiso + shurui


n, k = map(int, input().split())
t = []
d = []

for i in range(n):
    ti, di = map(int, input().split())
    t.append(ti)
    d.append(di)

max_manzoku = 0
for c in itertools.combinations(range(n), k):
    s = satisfaction(t, d, c)
    if s > max_manzoku:
        max_manzoku = s

print(max_manzoku)

Submission status: Time Limit Exceeded

Solution code:
from itertools import combinations
import numpy as np

n, k = map(int, input().split())
sl = [[int(i) for i in input().split()] for _ in range(n)]

sa = np.array(list(combinations(sl, k)))
c = sa.shape[0]
y = np.sum(sa[:,:,1], axis=1).reshape(c)
k = np.power(np.array([np.unique(l).shape[0] for l in sa[:,:,0]]), 2)
print(np.max(y + k))

Submission status: Time Limit Exceeded

Solution code:
import itertools
import collections

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

for i in list(itertools.combinations(list(range(N)), K)):
    num = 0
    ch = collections.defaultdict(int)
    for j in i:
        num += td[j][1]
        ch[td[j][0]] += 1
    ans = max(ans, num + len(ch)**2)

print(ans)

Submission status: Time Limit Exceeded

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  IN = list(map(int,input().split()))
  A.append(IN)

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
k = False
D = []
for i in range(K):
  while True:
    j += 1
    if j == N:
      j = 0
      k = True
    if A[j][1] != 0:
      if k:
        ans2 += A[j][1]
        D.append(A[j][0])
        A[j][1] = 0
        break
      if not A[j][0] in D:
        ans2 += A[j][1]
        D.append(A[j][0])
        A[j][1] = 0
        break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

Submission status: Time Limit Exceeded

Solution code:
IN = list(map(int,input().split()))
N = IN[0]
K = IN[1]

A = []

for i in range(N):
  IN = list(map(int,input().split()))
  A.append(IN)

A.sort(key=lambda x:-x[1])
ans1 = sum([a[1] for a in A[0:K]]) + len(set([a[0] for a in A[0:K]])) ** 2
ans2 = 0

j = -1
k = False
D = []
for i in range(K):
  while True:
    j += 1
    if j == N:
      j = 0
      k = True
    if A[j][1] != 0:
      if not A[j][0] in D or k:
        ans2 += A[j][1]
        D.append(A[j][0])
        A[j][1] = 0
        break
ans2 += len(set(D)) ** 2

print(max(ans1,ans2))

Submission status: Time Limit Exceeded

Solution code:
import itertools
def satis(sx):
    b = 0
    types = []
    for s in sx:
        b += s['d']
        n = s['n']
        if n not in types:
            types.append(n)
    return b + (len(types) ** 2)

n, k = map(int, input().split(' '))
sushis = []
for i in range(n):
    n, d = map(int, input().split(' '))
    sushis.append({'n': n, 'd': d})
print(max(map(satis, (itertools.combinations(sushis, k)))))

Submission status: Time Limit Exceeded

Solution code:
n, k = map(int, input().split())

t, d = [], []
for i in range(n):
    ti, di = map(int, input().split())
    t.append(ti)
    d.append(di)

#%%
# n, k = 5, 3
# t = [1, 1, 2, 2, 3]
# d = [9, 7, 6, 5, 1]


#%%
import itertools

#%%
ans = 0
for comb in list(itertools.combinations(range(n), k)):
    delicious = 0
    kinds = []
    kinds_point = 0

    for c in comb:
        delicious += d[c]
        if t[c] not in kinds:
            kinds.append(t[c])
            kinds_point +=1
    kinds_point = kinds_point**2

    satisfaction = delicious + kinds_point
    if satisfaction > ans:
        ans = satisfaction
print(ans)        

#%%

Submission status: Time Limit Exceeded

Solution code:
import itertools

N, K = map(int, input().split())

td = [[0, 0] for i in range(N)]
for i in range(N):
  td[i] = [int(i) for i in input().split()]

max = 0

c_list = list(itertools.combinations(td, K))

for i in range(len(c_list)):
  sum = 0
  types = []
  for c in c_list[i]:
    sum += c[1]
    types.append(c[0])
  
  num_types = len(set(types))

  score = sum + num_types * num_types

  if score > max:
    max = score

print(max)

Submission status: Time Limit Exceeded

Solution code:
import itertools

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

for i in list(itertools.combinations(td, K)):
    ch = []
    num = 0
    for t, d in i:
        num += d
        ch.append(t)
    num += len(set(ch)) ** 2
    ans = max(ans, num)
print(ans)

Submission status: Time Limit Exceeded

Solution code:
import numpy as np

N, K = map(int, input().split())
T = []
D = []
for i in range(N):
  t, d = map(int, input().split())
  T.append(t)
  D.append(d)

idx = np.argsort(D)[::-1]

res = 0
for x in range(1,N+1):
  used = {}
  m = 0
  k = 0
  y = 0
  for i in idx:
    if m == x:
      break
    if k == K:
      break
    t = T[i]
    d = D[i]
    if x - m == K - k:
      if t in used:
        continue
    if t not in used:
      used[t] = 1
      m += 1
    y += d
    k += 1
    res = max(res, m*m+y)
print(res)

Submission status: Accepted

Solution code:
from collections import defaultdict

N, K = map(int, input().split())
Items = [tuple(map(int, input().split())) for _ in range(N)]
Items.sort(reverse=True, key=lambda x: x[1])
Kinds = defaultdict(int)
ans = 0
k = 0
Candidates = []
for i, (t, d) in enumerate(Items):
    if i < K:
        ans += d
        if t not in Kinds:
            k += 1
        else:
            Candidates.append(d)
        Kinds[t] += 1
        if i == K - 1:
            ans += k ** 2
            a = ans
    else:
        if Kinds[t]:
            continue
        if not Candidates:
            break
        else:
            x = Candidates.pop()
            a -= x
            a -= k ** 2
            k += 1
            a += d
            a += k ** 2
            Kinds[t] -= 1
            ans = max(ans, a)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heapify, heappop, heappush
from collections import deque

n,k = map(int, input().split())
dtl = []
# t_dmax = [-1]*(n+1)
for _ in range(n):
    t,d = map(int, input().split())
    dtl.append((d,t))


dtl.sort(reverse=True)
dtq = deque(dtl)

dt_remq = deque([])
already_t = [False]*(n+1)
tll = [ [] for _ in range(n+1)]
cnt = 0
oisi = 0
kl = []
while dtq:
    d,t = dtq.popleft()
    if k <= cnt or already_t[t]: 
        dt_remq.append((d,t))
    else:
        already_t[t] = True
        kl.append((d,t))
        tll[t].append(d)
        cnt += 1
        oisi += d


rem = k-cnt
# if rem == 0:
#     ans = oisi + k*k
#     print(ans)
#     exit()

more_twos = []
for _ in range(rem):
    d,t = dt_remq.popleft()
    tll[t].append(d)
    oisi += d

ones = []
for tl in tll:
    if len(tl) == 1: ones.append(tl[0])

ones.sort()
onesq = deque(ones)

cuur_v = oisi + cnt*cnt
# print(cuur_v)
# print(kl,cnt)
while dt_remq and onesq:
    d,t = dt_remq.popleft()
    oisi_diff = d-onesq[0]
    bonus_diff = cnt*cnt-(cnt-1)*(cnt-1)
    if oisi_diff > bonus_diff:
        cuur_v += (oisi_diff-bonus_diff)
        cnt -= 1
        onesq.popleft()
    else:
        break

print(cuur_v)

Submission status: Accepted

Solution code:
from collections import Counter,defaultdict,deque
from heapq import heappop,heappush
from bisect import bisect_left,bisect_right 
import sys,math,itertools,fractions,pprint
sys.setrecursionlimit(10**8)
mod = 10**9+7
INF = float('inf')
def inp(): return int(sys.stdin.readline())
def inpl(): return list(map(int, sys.stdin.readline().split()))

n,k = inpl()
sushis = [] #[美味しさ,種類]
for _ in range(n):
    a,b = inpl()
    sushis.append([b,a])
sushis.sort(reverse = True)
sushis = deque(sushis)
cnt = [0] * (n+1)
kind = 0
ois_sum = 0
eats = []
for i in range(k):
    ois,knd = sushis.popleft()
    if cnt[knd] == 0:
        kind += 1
    cnt[knd] += 1
    ois_sum += ois
    eats.append([ois,knd])
eats.sort()
res = ois_sum + kind**2
ind = 0
while sushis and ind < k:
    while ind < k:
        eat_knd = eats[ind][1]
        if cnt[eat_knd] <= 1: 
            ind += 1
        else: break
    if ind >= k: break
    ois,knd = sushis.popleft()
    if cnt[knd]: continue
    eat_ois,eat_knd = eats[ind]
    cnt[knd] += 1; cnt[eat_knd] -= 1
    ois_sum -= eat_ois - ois
    kind += 1
    res = max(res, ois_sum + kind**2)
    ind += 1
print(res)

Submission status: Accepted

Solution code:
import sys
from collections import deque
from heapq import heappop, heappush

sys.setrecursionlimit(10 ** 7)
input = sys.stdin.readline
f_inf = float('inf')
mod = 10 ** 9 + 7


def resolve():
    n, k = map(int, input().split())
    TD = [list(map(int, input().split())) for _ in range(n)]
    TD.sort(key=lambda z: -z[1])
    TD_new = [[] for _ in range(n)]
    for t, d in TD:
        TD_new[t - 1].append([d, 0] if len(TD_new[t - 1]) else [d, 1])

    ALL = []
    for td in TD_new:
        ALL.extend(td)
    ALL.sort(reverse=True)
    que = deque(ALL)

    x = y = 0
    zero = []
    for _ in range(k):
        d, ex = que.popleft()
        x += d
        y += ex
        if ex == 0:
            heappush(zero, d)

    res = x + y ** 2
    while que:
        d, ex = que.popleft()
        if ex:
            if zero:
                x += d - heappop(zero)
                y += 1
                res = max(res, x + y ** 2)
            else:
                break

    print(res)


if __name__ == '__main__':
    resolve()

Submission status: Accepted

Solution code:
from heapq import heappush, heappop


N, K, *td = map(int, open(0).read().split())
sushis = [(d, t) for t, d in zip(*[iter(td)] * 2)]
sushis.sort(reverse=True)

pq = []
cnt = [0] * (N + 1)
base, x = 0, 0
for d, t in sushis[:K]:
    base += d
    if cnt[t] > 0:
        heappush(pq, (d, t))
    elif cnt[t] == 0:
        x += 1
    cnt[t] += 1
ans = base + x * x

for d, t in sushis[K:]:
    if not pq:
        break
    if cnt[t] > 0:
        continue

    _d, _t = heappop(pq)
    cnt[_t] -= 1
    cnt[t] += 1
    x += 1
    base += d - _d
    ans = max(ans, base + x * x)

print(ans)

Submission status: Accepted

Solution code:
import sys
read = sys.stdin.readline
n, k = map(int, read().split())
# 種類数を増やすときは、各種類で最もでかいのを食べる
# 種類数最大降順配列、他降順配列みたいなの用いしtえ
# 食べる出井数を全探索する
# 累積和
sushi = [[] for _ in range(n)]

for i in range(n):
    t, d = map(int, read().split())
    t -= 1
    sushi[t].append(d)

max_sushi = []
other_sushi = []
for i in range(n):
    if len(sushi[i]) == 0:
        continue
    sushi[i].sort()
    max_sushi.append(sushi[i][-1])
    for j in range(len(sushi[i]) - 1):
        other_sushi.append(sushi[i][j])
max_sushi.sort(reverse=True)
other_sushi.sort(reverse=True)

for i in range(len(max_sushi) - 1):
    max_sushi[i + 1] += max_sushi[i]
for i in range(len(other_sushi) - 1):
    other_sushi[i + 1] += other_sushi[i]

a = len(max_sushi)
b = len(other_sushi)
ans = 0
for x in range(1, k + 1):
    if x > a or k - x > b:
        continue
    other = 0 if k - x - 1 < 0 else other_sushi[k - x - 1]
    ans = max(ans, x * x + max_sushi[x - 1] + other)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heapify, heappop
from sys import stdin
def input():
    return stdin.readline().strip()

def main():
    n, k = map(int, input().split())

    dt = [(0, 0)] * n
    for i in range(n):
        t, d = map(int, input().split())
        dt[i] = (d, t-1)

    dt.sort(key=lambda x: x[0], reverse=True)

    # select k biggest points
    base = 0
    variety = [0] * n
    bonus = 0
    heap = []
    for d, t in dt[:k]:
        base += d
        if variety[t] == 0:
            bonus += 1
        variety[t] += 1
        heap.append((d, t))
    heapify(heap)

    # select exchange candidates
    new_dt = []
    for d, t in dt[k:]:
        if variety[t] == 0:
            variety[t] = 1
            new_dt.append((d, t))

    # exchange
    ans = base + bonus * bonus
    for new_d, new_t in new_dt:
        d, t = heappop(heap)
        while len(heap) > 0 and variety[t] == 1:
            d, t = heappop(heap)
        if len(heap) == 0:
            break
        variety[t] -= 1

        base -= d - new_d
        bonus += 1
        if ans < base + bonus * bonus:
            ans = base + bonus * bonus

    print(ans)

main()