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: Accepted

Solution code:
def takeSecond(elem):
    return elem[1]

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

sushi = [list(map(int,input().split())) for i in range(n)]

sushi.sort(key=takeSecond,reverse=True)
'''
selected = sushi[:k]
kindset = set([item[0] for item in selected])

startKindNum = len(kindset)

score = startKindNum*startKindNum + sum([item[1] for item in selected])
maxscore = score
front_pointer = k - 1
pointer = k
for i in range(startKindNum+1,min([k+1,n+1])):
    #remove the smallest one that does not decrease kind
    for p in range(front_pointer,-1,-1):
        item = selected[p]
        flag = False
        for idx in range(p):
            if(selected[idx][0]==item[0]):
                score -= item[1]
                selected.remove(item)
                front_pointer = p - 1
                flag = True
                break
        if(flag):
            break

    #add the largest one that add kind
    for j in range(pointer,n):
        if(sushi[j][0] not in kindset):
            selected.append(sushi[j])
            kindset.add(sushi[j][0])
            k += 1
            pointer = j + 1            
            score += 2*startKindNum + 1 + sushi[j][1]
            if(score>maxscore):
                maxscore = score
            startKindNum += 1
            break
'''
kindset = set([])
score = 0
extras = []
for item in sushi[:k]:
    if item[0] in kindset:
        extras.append(item[1])
    else:
        kindset.add(item[0])
    score += item[1]
score += len(kindset)**2
maxscore = score
for item in sushi[k:]:
    if not  extras:
        break
    if item[0] in kindset:
        continue
    score += (item[1] - extras.pop()) + 2*len(kindset) + 1
    kindset.add(item[0])
    #print(score)
    if score > maxscore:
        maxscore = score


print(maxscore)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi.sort(key=lambda x: x[1], reverse=True)

score = 0

remains = []
kinds = set()
for t, d in sushi[:K]:
    if t in kinds:
        remains.append(d)
    kinds.add(t)
    score += d
score += len(kinds)**2

test = score
for t, d in sushi[K:]:
    if not remains:
        break
    if t in kinds:
        continue
    kinds.add(t)
    rm = remains.pop()
    test = test + (d - rm) + (len(kinds)*2 - 1)
    score = max(score, test)

print(score)

Submission status: Accepted

Solution code:
import heapq as hq
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    hq.heappush(stack,[d,t])
ans = kiso + len(neta)**2
for t,d in t_d[K:]:
    if neta[t] > 0 : continue
    while(stack):
        d0,t0 = hq.heappop(stack)
        if neta[t0] > 1:
            kiso += d - d0
            neta[t0] -= 1
            neta[t] += 1
            ans = max(ans, kiso+len(neta)**2)
            hq.heappush(stack,[d0,t0])
            hq.heappop(stack)
            break
print(ans)

Submission status: Accepted

Solution code:
import heapq as hq
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    hq.heappush(stack,[d,t])
ans = kiso + len(neta)**2
for t,d in t_d[K:]:
    if neta[t] > 0 : continue
    while(stack):
        d0,t0 = hq.heappop(stack)
        if neta[t0] > 1:
            kiso += d - d0
            neta[t0] -= 1
            neta[t] += 1
            ans = max(ans, kiso+len(neta)**2)
            break
print(ans)

Submission status: Accepted

Solution code:
from collections import Counter
N,K = map(int,input().split())
t_d = []
for _ in range(N):
    t,d = map(int,input().split())
    t_d.append([t,d])
t_d = sorted(t_d, key=lambda x:-x[1])
kiso = 0
neta = Counter()
stack = []
for t,d in t_d[:K]:
    kiso += d
    neta[t] += 1
    stack.append((t,d))
ans = kiso + len(neta)**2
tmp = ans
for t,d in t_d[K:]:
    if t in neta : continue
    while(stack):
        t0,d0 = stack.pop()
        if neta[t0] > 1:
            tmp += d - d0 + 2*len(neta) + 1
            ans = max(ans,tmp)
            neta[t0] -= 1
            neta[t] += 1
            break
    else:
        break
print(ans)

Submission status: Accepted

Solution code:
def solve(n, k, ts, ds, dt):
    firsts = set()
    seconds = []

    dt.sort(reverse=True)

    count = 0
    t_sum = 0
    d_sum = 0
    score_max = 0
    for d, t in dt:

        if count == k:
            # かぶり一切なし
            if not seconds:
                break

            # 今までにいれたことある種類はパス
            if t in firsts:
                continue

            d_sum -= seconds[-1]
            seconds.pop()
            count -= 1

        if t not in firsts:
            firsts.add(t)
            t_sum += 1
        else:
            seconds.append(d)

        d_sum += d
        score = d_sum + (t_sum * t_sum)
        # print(d_sum, ',', t_sum, '=', score)

        if score_max < score:
            score_max = score

        count += 1

    return score_max


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

ts = [0] * n
ds = [0] * n
td = [(0, 0)] * n
for i in range(n):
    t, d = map(int, input().split())
    ts[i] = t
    ds[i] = d
    td[i] = (d, t)

print(solve(n, k, ts, ds, td))

Submission status: Accepted

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 i in range(k):
    d,t = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)

x = sum(X)
y = sum(Y)
v = len(Y)
ans = x + y + v**2
for i in range(k, n):
    if not X:
        break
    d, t = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v**2)

print(ans)

Submission status: Accepted

Solution code:
# 入力
n, k = map(int, input().split())
A = [] # 寿司のネタの美味しさと種類を入れるリスト
for _ in range(n):
    A.append(list(map(int, input().split())))

# 美味しさ順にソートする
A.sort(key=lambda x: x[1], reverse=True)

# 美味しさで貪欲にとる
X = [] # 同じ種類で他に美味しさの大きい寿司が存在する
Y = [] # 取った寿司の中の各々の種類で美味しさが最も大きい
st = set() # 種類
for i in range(k):
    t, d = A[i]
    if t in st:
        X.append(d)
    else:
        Y.append(d)
        st.add(t)

# 美味しさポイントを求める
x = sum(X)
y = sum(Y)
v = len(Y) # 種類数
ans = x + y + v ** 2
for i in range(k, n):
    if not X:
        break
    t, d = A[i]
    if t in st:
        continue
    st.add(t)
    x -= X.pop()
    y += d
    v += 1
    ans = max(ans, x + y + v ** 2)

print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict
import heapq


def get_first_try(K, ts, ds):
    eats = []
    non_eats = []
    for i in range(len(ts)):
        t = ts[i]
        d = ds[i]
        heapq.heappush(eats, (d, t))
        if len(eats) > K:
            non_eats.append(heapq.heappop(eats))
    return [(t, d) for d,t in eats], [(t, d) for d,t in non_eats]

def count_kinds(eats):
    kinds = defaultdict(int)
    for e in eats:
        kinds[e[0]] += 1
    return kinds

def calc_satisfaction(eats, eat_kinds):
    s = 0
    for e in eats:
        s += e[1]
    return s + calc_bonus(eat_kinds)


def calc_bonus(kinds):
    x = 0
    for k in kinds:
        if kinds[k] > 0:
            x += 1
    return x**2


def solve(N, K, ts, ds):
    eats, non_eats = get_first_try(K, ts, ds)
    eat_kinds = count_kinds(eats)
    satisfaction = calc_satisfaction(eats, eat_kinds)
    sorted_eats = sorted(eats, key=lambda x: x[1])
    sorted_non_eats = sorted(non_eats, key=lambda x: x[1], reverse=True)
    n_kinds = len(eat_kinds)

    max_satisfaction = satisfaction

    i = 0
    j = 0
    while i < len(sorted_eats) and j < len(sorted_non_eats):
        t, d = sorted_eats[i]
        if eat_kinds[t] == 1:
            i += 1
            continue
        t2, d2 = sorted_non_eats[j]
        if eat_kinds[t2] > 0:
            j += 1
            continue
        prev_bonus = n_kinds ** 2
        eat_kinds[t] -= 1
        eat_kinds[t2] += 1
        n_kinds += 1
        current_bonus = n_kinds ** 2
        satisfaction = satisfaction - d + d2 - prev_bonus + current_bonus
        if max_satisfaction < satisfaction:
            max_satisfaction = satisfaction
        i += 1
        j += 1

    return max_satisfaction


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

ts = []
ds = []
for i in range(N):
    t, d = list(map(int, input().split(' ')))
    ts.append(t)
    ds.append(d)

print(solve(N, K, ts, ds))

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
sushi = []
for i in range(n):
    tmp = list(map(int, input().split()))
    sushi.append(tmp)
eaten = set() # 食べた種類を記憶(集合型)
chohuku = [] # 同じ種類で二回以上食べた寿司のスコアを記録
scores = [] # Sのスコアのリスト

sushi.sort(key=lambda x:x[1])
sushi.reverse() # sushiをおいしさの降順にソート

tmp = 0 # Sのおいしさの合計を保持
for i in range(k): # sushiを美味しい順にk個食べる
    if sushi[i][0] in eaten: # 同じ種類の寿司を食べていた場合
        chohuku.append(sushi[i][1])
    else:
        eaten.add(sushi[i][0]) # 食べた寿司の種類を記憶
    tmp += sushi[i][1] # 食べた寿司のおいしさをtmpに追加
chohuku.sort()
scores.append(tmp+len(eaten)**2) # 今のSのスコアを記憶

for i in range(k,n):
    if not sushi[i][0] in eaten and chohuku != []:
        tmp -= chohuku.pop(0)
        eaten.add(sushi[i][0])
        tmp += sushi[i][1]
        scores.append(tmp+len(eaten)**2)
        
print(max(scores))

Submission status: Accepted

Solution code:
from collections import defaultdict

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

sushi = sorted(sushi, key = lambda x: x[1], reverse = True)
types = set()
max_vals = [0]
other_vals = [0]

for typ, val in sushi:
    if typ not in types:
        max_vals.append(max_vals[-1] + val)
        types.add(typ)
    else:
        other_vals.append(other_vals[-1] + val)

max_score = 0
for k in range(1, K + 1):
    if K - k < len(other_vals) and k < len(max_vals):
        score = max_vals[k] + other_vals[K - k] + k ** 2
        max_score = max(max_score, score)

print(max_score)

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi = sorted(sushi, key=lambda x: x[1], reverse=True)
typ = set()
lstA = []
lstB = []
for item in sushi:
    if item[0] in typ:
        item[0] = 0
        lstB.append(item[1])
    else:
        typ.add(item[0])
        item[0] = 1
        lstA.append(item[1])
ans = 0
neta = 0
for i in range(K):
    ans += sushi[i][1]
    if sushi[i][0] == 1:
        neta += 1
ans += neta ** 2
C = sum(lstA[:neta])
D = sum(lstB[:K - neta])
for i in range(neta + 1, min(K, len(lstA)) + 1):
    C += lstA[i - 1]
    D -= lstB[K - i]
    tmp = C + D + i ** 2
    ans = max(ans, tmp)
print(ans)

Submission status: Accepted

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

n, k = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(n)]
sushi.sort(key = lambda x: x[1], reverse = True)    # sort by d

mp_max = 0
op = 0
taberu_ind = set()
pop_cands = []
t_taberu = set()
# greedy for delicious point
for i in range(k):
    op += sushi[i][1]

    t = sushi[i][0]
    if t_taberu.isdisjoint({t}):    # new type
        t_taberu.add(t)
    else:
        pop_cands.append(i)

    taberu_ind.add(i)
mp_max = op + len(t_taberu)**2

# reduce dericious, gain variety
i += 1
j = i
for i in range(i, n):
    if len(pop_cands) == 0:
        break

    pop_ind = pop_cands.pop()
    taberu_ind.remove(pop_ind)

    for j in range(j, n):
        t = sushi[j][0]
        if t_taberu.isdisjoint({t}):
            t_taberu.add(t)
            break
    else:
        break
    
    # calc points
    op -= sushi[pop_ind][1]
    op += sushi[j][1]
    mp = op + len(t_taberu)**2
    mp_max = max(mp_max, mp)

    taberu_ind.add(i)    # not necessary actually
print(mp_max)

Submission status: Accepted

Solution code:
N, K = map( int, input().split())
T = [(0,0)]*N
for i in range(N):
    t, d = map( int, input().split())
    T[i] = (d,t)
T.sort( key= None, reverse = True)
C = T[:K]
AC = T[K:]
kiso = 0
s = 0
S = [0]*(N+1)
for i in range(K):
    d, t = T[i]
    kiso += d
    if S[t] == 0:
        S[t] = 1
        s += 1
    else:
        S[t] += 1
ans = kiso + s**2
C.sort()
r = 0
 
for i in range(K):
    d, t = C[i]
    if S[t] >= 2:
        while r < N-K:
            if S[ AC[r][1]] == 0:
                S[t] -= 1
                S[ AC[r][1]] += 1
                s += 1
                kiso -= d
                kiso += AC[r][0]
                r += 1
                break
            r += 1
    ans = max( ans, kiso + s**2)
print( ans)

Submission status: Accepted

Solution code:
#ABC116D-VariousSushi.py

from sys import stdin
#from numpy import sum

input = stdin.readline

N, K = list(map(int, input().split()))
List1= []
kind = set()
point1, point2 = 0, 0

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

INPUT.sort(key = lambda x: -x[1])

for ti, di in INPUT[:K]:
    point1 += di
    if ti in kind:
        List1.append(di)
    kind.add(ti)

point2 = len(kind)

Score = point1 + point2**2
Test = Score

List1POP = List1.pop

for ti, di in INPUT[K:]:
    if 0 >= len(List1):
        break
    if ti in kind:
        continue
    Test = Test - List1POP() + di + 2 * point2 + 1
    point2 += 1
    kind.add(ti)
    Score = max(Score, Test)

print(Score)

Submission status: Accepted

Solution code:
# -*- coding: utf-8 -*-
from sys import stdin
from collections import Counter
import numpy as np

s_in = lambda: stdin.readline()[:-1]  # s = s_in()
d_in = lambda: int(stdin.readline())  # N = d_in()
ds_in = lambda: list(map(int, stdin.readline().split()))  # List = ds_in()


N, K = ds_in()
t_list = []
d_list = []
for _ in range(N):
    t, d = ds_in()
    t_list.append(t)
    d_list.append(d)
np_t = np.array(t_list)  # sushi kind
np_d = np.array(d_list)  # sushi value

sort_idx = np.argsort(np_d)[::-1][:K]  # top K value sushi
top_sushi = np_t[sort_idx]
top_score = np_d[sort_idx]
kind = len(set(top_sushi))
point = top_score.sum() + kind**2  # temp max point
point_list = [point]
sushi_counter = Counter(top_sushi)

sushi_set = {k: [] for k in set(t_list)}
for i, j in zip(t_list, d_list):
    sushi_set[i].append(j)
max_sushi = {k: max(sushi_set[k]) for k in set(t_list)}

for s in set(top_sushi):
    max_sushi.pop(s)  # temp max point
max_values = list(max_sushi.values())
max_values.sort(reverse=True)
if len(max_values) > 0:
    for sushi, score in zip(top_sushi[::-1], top_score[::-1]):
        if sushi_counter[sushi] > 1:
            sushi_counter[sushi] -= 1
            kind += 1
            point = point - score + max_values.pop(0) + 2*kind - 1
            point_list.append(point)
            if len(max_values) == 0:
                break

print(max(point_list))

Submission status: Accepted

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

td=[]

for i in range(n):
    a=[]
    a=[int(j) for j in input().split()]
    td.append(a)

from operator import itemgetter
td.sort(key=itemgetter(1))  
td.reverse()
td.sort(key=itemgetter(0))

s=0
l=[]
for i in range(n):
    if s!=td[i][0]:
        s=td[i][0]
        l.append([td[i][1],1])
    else:
        l.append([td[i][1],0])

l.sort(key=itemgetter(0))
l.reverse()

op=0
sp=0
ans=[]
sl=[]
l1=[]

for i in range(n):
    if i<k:
        op+=l[i][0]
        sp+=l[i][1]
        if l[i][1]==0:
            sl.append(l[i][0])
    else:
        if l[i][1]==1:
            l1.append(l[i][0])


ans=[op+sp**2]
sl.sort()
l1.sort()
l1.reverse()

for i in range(min(len(sl),len(l1))):
    ans.append(ans[-1]-sl[i]+l1[i]+sp*2+1)
    sp+=1

print(max(ans))

Submission status: Accepted

Solution code:
from collections import deque

if __name__ == '__main__':
    n, k = map(int, input().split())
    sushis = [tuple(int(s) for s in input().split()) for _ in range(n)]
    sushis.sort(key=lambda t: t[1], reverse=True)

    used_set = set()
    replace_stack = deque()
    add_queue = deque()
    ans = 0
    for i in range(k):
        t, d = sushis[i]
        if t in used_set:
            replace_stack.appendleft(d)
        else:
            used_set.add(t)
        ans += d
    num = len(used_set)
    ans += num * num

    for i in range(k, n):
        t, d = sushis[i]
        if t not in used_set:
            add_queue.append(d)
            used_set.add(t)

    tmp = ans
    for x, r, a in zip(range(num, k + 1), replace_stack, add_queue):
        tmp = tmp + 2 * x + 1 - r + a
        ans = max(ans, tmp)

    print(ans)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import sys
input = sys.stdin.readline

def main():
  N, K = map(int, input().split())
  sushis = [tuple(map(int, input().split())) for _ in range(N)]
  sushis.sort(key=lambda x:x[1], reverse=True)
  curr = sushis[:K]
  leave = sushis[K:]
  base = lambda li: sum([l[1] for l in li])
  types_max = len(set([c[0] for c in sushis]))
  types = set()
  dups = []

  for c in curr:
    if c[0] in types:
      dups += [c]
    else:
      types.add(c[0])
  types = set(types)

  ps = base(curr) + len(types)**2
  di = {}
  di[len(types)] = ps
  n_types = len(types) + 1
  for l in leave:
    if len(types) >= types_max:
      break
    if len(dups) == 0:
      break
    if not l[0] in types:
      dup = dups.pop()
      di[n_types] = ps - dup[1] + l[1] + n_types**2 - (n_types-1)**2
      ps = di[n_types]
      n_types += 1
      types.add(l[0])

  print(sorted(di.items(), key=lambda x:x[1])[-1][1])


if __name__=='__main__':
  main()

Submission status: Accepted

Solution code:
import sys

def point(sushi, n, k):
  sushi = sorted(sushi, key=lambda x:-x[1])
  
  ate_t = set()
  dpoint = 0
  dup = []
  for s in sushi[:k]:
    t = s[0]
    d = s[1]
    if t in ate_t:
      dup.append(d)
    ate_t.add(t)
    dpoint += d
  
  ans = len(ate_t)**2 + dpoint
  
  for s in sushi[k:]:
    t = s[0]
    d = s[1]
    if len(dup) <= 0:
      break
    if t in ate_t:
      continue
    ate_t.add(t)
    rm = dup.pop()
    dpoint = dpoint - rm + d
    res = len(ate_t)**2 + dpoint
    ans = max(ans, res)
  
  return ans

def main():
  input = sys.stdin.readline
  n, k = map(int, input().split())
  sushi = []
  for _ in range(n):
    t, d = map(int, input().split())
    sushi.append([t,d])
  print(point(sushi, n, k))


if __name__ == '__main__':
  main()

Submission status: Accepted

Solution code:
class Sushi(object):
    def __init__(self, sushi_type, delicious):
        self.t = sushi_type
        self.d = delicious


def main() -> None:
    n, k = map(int, input().split())
    menu = [[int(x) for x in input().split()] for _ in range(n)]
    selected = []
    unselected = []
    cnt = 0
    p_satisfied = 0
    selected_type = set()

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

    for i in range(n):
        sushi = Sushi(menu[i][0], menu[i][1])
        if sushi.t not in selected_type and cnt < k:
            selected_type.add(sushi.t)
            cnt += 1
            selected += [menu[i]]
        else:
            unselected += [menu[i]]

    unselected.sort(key=lambda x: x[1])

    while cnt < k and len(unselected) > 0:
        p_satisfied += unselected[-1][1]
        unselected.pop()
        cnt += 1

    while True:
        t_cnt = len(selected)
        if len(selected) == 0 or len(unselected) == 0:
            break
        if t_cnt ** 2 + selected[-1][1] > (t_cnt - 1) ** 2 + unselected[-1][1]:
            break
        p_satisfied += unselected[-1][1]
        unselected.pop()
        selected.pop()

    for v in selected:
        p_satisfied += v[1]

    print(p_satisfied + len(selected) ** 2)


if __name__ == "__main__":
    main()

Submission status: Accepted

Solution code:
import sys

N,K=map(int, sys.stdin.readline().split()) 
T=[ map(int, sys.stdin.readline().split()) for i in range(N) ]
T=sorted(T, key=lambda x:(x[0],-x[1]))

KI=[]	#KIND
A=[]	#Nokori

x=None
for t,d in T:
	if t != x:
		KI.append(d)
		x=t
	else:
		A.append(d)

KI.sort(reverse=True)
A.sort(reverse=True)

KI_S=[0]		#KI Ruiseki sum
A_S=[0]		#A Ruiseki sum

for x in KI: KI_S.append( KI_S[-1]+x )
for x in A:  A_S.append( A_S[-1]+x )


la=len(A_S)
lk=len(KI_S)

ans=0
for k in range(1,K+1):
	if K-k<la and k<lk:
		ans=max (ans , KI_S[k]+k**2+A_S[K-k] )

print ans

Submission status: Accepted

Solution code:
a = list(map(lambda x: int(x), input().split()))
N = a[0]
K = a[1]
sushis = []
for i in range(0, N):
  sushis.append(list(map(lambda x: int(x), input().split())))

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

netas = set()
duplicates = []

choices = sushis[:K]

for choice in choices:
  if choice[0] in netas:
    duplicates.append(choice[1])
  else:
    netas.add(choice[0])

p = sum(map(lambda x: x[1], choices)) + len(netas)**2
answer = p

for sushi in sushis[K:]:
  if sushi[0] not in netas and duplicates:
    netas.add(sushi[0])
    duplicate = duplicates.pop()
    p += (sushi[1] - duplicate) + 2*len(netas) - 1
    answer = max(answer, p)

print(answer)

Submission status: Accepted

Solution code:
from collections import Counter as ct
n, k = map(int, input().split())
td = [list(map(int, input().split())) for i in range(n)]

kind = len(ct(list(map(lambda x: x[0], td))))
std = sorted(td, key=lambda x: x[1], reverse=True)
dup = []
eat_k = set()
ans = 0
for i in range(k):
    ans += std[i][1]
    if std[i][0] in eat_k:
        dup.append(std[i])
    eat_k.add(std[i][0])
ans += len(eat_k)** 2
tmpa = ans
for i in range(k, n):
    if len(dup) == 0:
        break
    if std[i][0] not in eat_k:
        t = dup.pop()
        tmpa -= t[1]
        tmpa += std[i][1]
        tmpa -= len(eat_k)** 2
        eat_k.add(std[i][0])
        tmpa += len(eat_k)** 2
        ans = max(tmpa, ans)
print(ans)

Submission status: Accepted

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

for i in range(N):
    t, d = map(int, input().split())
    menu_.append((t, d))

menu = sorted(menu_, key = lambda x:x[1], reverse=True)
neta = set()
duplicates = []

for m in menu[:K]:
    if m[0] in neta:
        duplicates.append(m[1])
    else:
        neta.add(m[0])

p = sum([d for (t, d) in menu[:K]]) + len(neta)**2
ans = p

for m in menu[K:]:
    if m[0] not in neta and duplicates:
        neta.add(m[0])
        duplicate = duplicates.pop()
        
        p += (m[1] - duplicate) + 2 * len(neta) - 1 
        #print(ans, p)
        ans = max(ans, p)

print(ans)

Submission status: Accepted

Solution code:
from operator import itemgetter
from functools import reduce
N, K = map(int, input().split())
t, d = zip(*(map(int, input().split()) for _ in range(N)))
fn = itemgetter(0)
fd = itemgetter(1)
ss = sorted(zip(t, d), key=fd, reverse=True)
netas = set()
duplicates = []
for s in ss[:K]:
  if s[0] in netas:
    duplicates.append(s[1])
  else:
    netas.add(s[0])
p = sum(map(fd, ss[:K])) + len(netas)**2
ans = p
for s in ss[K:]:
  if s[0] not in netas and duplicates:
    netas.add(s[0])
    duplicate = duplicates.pop()
    p += (s[1] - duplicate) + 2 * len(netas) - 1
    ans = max(ans, p)
print(ans)

Submission status: Accepted

Solution code:
from heapq import heappop,heappush
n,k = (int(i) for i in input().split())
x,y,z,q = [],set(),[0]*n,[]
for _ in range(n):
	t,d = (int(i) for i in input().split())
	x.append((d,t-1))
x,p = sorted(x)[::-1],0
for i in range(k):
	p+=x[i][0]
	z[x[i][1]]+=1
	y.add(x[i][1])
	heappush(q,x[i])
ans,m,l,F = p+len(y)**2,k,len(y),1
while F:
	F = 0
	while q:
		d,t = heappop(q)
		if z[t]!=1:
			f,F,z[t] = 1,1,z[t]-1
			for j in range(m,n):
				if z[x[j][1]]==0:
					z[x[j][1]],m,f,l = 1,j,0,l+1
					p = p-d+x[j][0]
					break
			if f: F = 0
			else: ans = max(ans,p+l**2)
			break
print(ans)

Submission status: Accepted

Solution code:
import heapq


def func(n, k, td):
    neta = [0 for _ in range(n)]
    td.sort()
    td.reverse()
    _heapq = []
    _sum = 0
    types = 0
    for i in range(k):
        _sum += td[i][0]
        if neta[td[i][1]] == 0:
            types += 1
        if neta[td[i][1]] > 0:  # 被っているものだけ後で削除する可能性あり
            heapq.heappush(_heapq, td[i])
        neta[td[i][1]] += 1
    _max = _sum + types**2
    # print("td:{} heapq:{}".format(td, _heapq))
    for i in range(k, n):
        # neta[td[i][1]] != 0なら,
        # td[i] < td[j](j<i)なのでtd[i]は無視
        if len(_heapq) > 0 and neta[td[i][1]] == 0:  # 種類が増えるとき
            tmp = heapq.heappop(_heapq)
            _sum -= tmp[0]
            _sum += td[i][0]
            neta[tmp[1]] -= 1
            types += 1
            neta[td[i][1]] += 1
        # print("td:{} heapq:{}".format(td, _heapq))
        # print("_sum:{} typed:{}".format(_sum, types))
        _max = max(_max, _sum + types**2)
    return str(_max)


n, k = [int(i) for i in input().split()]
td = [[0, 0] for _ in range(n)]
for i in range(n):
    td[i][1], td[i][0] = [int(i) for i in input().split()]
    td[i][1] -= 1
# print("print(func({},{},{}))".format(n, k, td))
print(func(n, k, td))
"""
print(func(5,3,[[9, 0], [7, 0], [6, 1], [5, 1], [1, 2]]))
print(func(7,4,[[1, 0], [1, 1], [1, 2], [6, 3], [5, 3], [5, 3], [5, 3]]))
print(func(6,5,[[1000000000, 4], [990000000, 1], [980000000, 2], [970000000, 5], [960000000, 5], [950000000, 3]]))
"""

Submission status: Accepted

Solution code:
N,K =map(int,input().split())
sushi = [[int(i)for i in input().split()]for _ in range(N)]
sushi =sorted(sushi,key=lambda x:x[1],reverse=True)
neta =set()
ans = 0
daburi =[]
for i in range(K):
    if sushi[i][0] in neta:
        daburi.insert(0,i)
    neta.add(sushi[i][0])
    ans += sushi[i][1]
n = len(neta)
ans += n*n
sta = K
anslist=[ans]
s=False
while True:
    n =len(neta)
    if daburi == [] or s:
        break
    outsushi = daburi.pop(0)
    for i in range(sta,N):
        if not sushi[i][0] in neta:
            ans -= sushi[outsushi][1]
            ans += sushi[i][1]+2*n+1
            anslist.append(ans)
            neta.add(sushi[i][0])
            sta = i +1
            break
        elif i == N-1:
            s =True

print(max(anslist))

Submission status: Accepted

Solution code:
N, K = map(int, input().split())
dt = []
for i in range(N):
	t, d = map(int, input().split())
	dt.append((d, t))
dt.sort()
dt.reverse()
dic = {}
can_delete = []
dp = 0
x = 0
for i in range(K):
	d, t = dt[i]
	dp += d
	if t in dic:
		can_delete.append(d)
	else:
		dic[t] = d
		x += 1
can_delete.sort()
can_delete.reverse()
ans = dp + x * x
v = ans
for i in range(K, N):
	if len(can_delete) == 0:
		break
	d, t = dt[i]
	if t in dic:
		continue
	dic[t] = d
	x += 1
	v = v - can_delete.pop() + d + 2 * x - 1
	if v > ans:
		ans = v
print(ans)

Submission status: Accepted

Solution code:
import heapq
N, K = map(int, input().split())
dt = []
for i in range(N):
	t, d = map(int, input().split())
	dt.append((d, t))
dt.sort()
dt.reverse()
dic = {}
can_delete = []
dp = 0
x = 0
for i in range(K):
	d, t = dt[i]
	dp += d
	if t in dic:
		heapq.heappush(can_delete, d)
	else:
		dic[t] = d
		x += 1
ans = dp + x * x
v = ans
for i in range(K, N):
	if len(can_delete) == 0:
		break
	d, t = dt[i]
	if t in dic:
		continue
	dic[t] = d
	x += 1
	v = v - heapq.heappop(can_delete) + d + 2 * x - 1
	if v > ans:
		ans = v
print(ans)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
h = []
for i in range(n):
    t, d = map(int, input().split())
    h.append((t,d))
h.sort(key=lambda x: -x[1])

taste = 0
typen = 0
typecnt = {}
bad = []


for i in range(k):
    typecnt[h[i][0]] = typecnt.get(h[i][0], 0) + 1
    if(typecnt[h[i][0]] == 1):
        typen += 1
    else:
        bad.append(h[i][1])
    taste += h[i][1]

# print(h)
# print(typecnt)
# print(bad)

ans = taste + typen**2

for i in range(k, n):
    if(len(bad) == 0):
        break
    cnt = typecnt.get(h[i][0], 0)
    if(cnt != 0):
        continue
    else:
        typecnt[h[i][0]] = cnt + 1
    typen += 1
    taste -= bad.pop()
    taste += h[i][1]
    ans = max(ans, taste + typen**2)

print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict

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

items = []
for _ in range(n):
    ti, di = map(int, input().split())
    items.append((di, ti))

items.sort(reverse=True)

topics = defaultdict(int)
for _, ti in items[:k]:
    topics[ti] += 1

best_res = sum(v for v, _ in items[:k]) + len(topics) ** 2
curr_res = sum(v for v, _ in items[:k])

i = k - 1
j = k
while i >= 0 and j < n:
    di, ti = items[i]
    if topics[ti] < 2:
        i -= 1
        continue
    dj, tj = items[j]
    if tj in topics:
        j += 1
        continue
    if topics[ti] > 1:
        topics[ti] -= 1
    else:
        del topics[ti]
    topics[tj] += 1
    curr_res -= di
    curr_res += dj
    best_res = max(curr_res + len(topics) ** 2, best_res)
    i -= 1
        
    

print(best_res)

Submission status: Accepted

Solution code:
#!/usr/bin/env python3
N, K = map(int, input().split())
sushi = [list(map(int, input().split())) for _ in range(N)]
sushi.sort(key=lambda x: x[1], reverse=True)

point = 0
stack = []
kind = set()

for t, d in sushi[:K]:
    point += d
    if t in kind:
        stack.append(d)
    else:
        kind.add(t)

point += len(kind) ** 2
ans = point

for t, d in sushi[K:]:
    if not stack:
        break
    if t not in kind:
        point += d - stack.pop() + 2 * len(kind) + 1
        kind.add(t)
        ans = max(ans, point)

print(ans)