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:
from collections import defaultdict
from itertools import accumulate
import sys
input = sys.stdin.readline

def main():
    N, K = map(int, input().split())
    d = defaultdict(list)
    for _ in range(N):
        t, v = map(int, input().split())
        d[t].append(v)

    other = []
    vmax = []
    for nums in d.values():
        nums.sort(reverse=True)
        for i, j in enumerate(nums):
            if i == 0:
                vmax.append(j)
            else:
                other.append(j)

    vmax.sort(reverse=True)
    other.sort(reverse=True)
    sum_v = [0] + list(accumulate(vmax))
    sum_o = [0] + list(accumulate(other))
    ans = 0
    for i in range(K,0,-1):
        neta = min(i, len(vmax))
        base = neta * neta
        tmp = base + sum_v[neta]
        if K - neta > len(other):
            break
        tmp += sum_o[K-neta]
        ans = max(ans, tmp)
        


    print(ans)

    
    




if __name__ == "__main__":
    main()

Submission status: Accepted

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

import sys

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10 ** 9)
INF = 10 ** 18
MOD = 10 ** 9 + 7

N, K = MAP()
DT = []
tmp = [[] for i in range(N)]
for i in range(N):
    t, d = MAP()
    DT.append((d, t))
DT.sort(reverse=1)

A = []
D = {}
dcnt = 0
# ポイント上位K個をまず取る
for i in range(K):
    d, t = DT[i]
    dcnt += d
    # 種類tが1つ目なら、確定で残す
    if t not in D:
        D[t] = d
    # 2つ目以降なら、最大のものを残して他は交換候補へ
    else:
        if D[t] < d:
            D[t], d = d, D[t]
        A.append(d)
A.sort(reverse=1)
tcnt = len(D)

ans = dcnt + tcnt**2
for i in range(K, N):
    d2, t2 = DT[i]
    # 既に持っている種類ならスキップ
    if t2 in D:
        continue
    D[t2] = d2
    # 全部1種類ずつ(交換候補がもうない)なら、もう種類数を増やせないので終了
    if not A:
        break
    d = A.pop()
    # この候補と交換してみる
    dcnt -= d
    dcnt += d2
    tcnt += 1
    ans = max(ans, dcnt + tcnt**2)
print(ans)

Submission status: Accepted

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

import sys
from heapq import heapify, heappop

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10 ** 9)
INF = 10 ** 18
MOD = 10 ** 9 + 7

N, K = MAP()
DT = []
tmp = [[] for i in range(N)]
for i in range(N):
    t, d = MAP()
    DT.append((d, t))
DT.sort(reverse=1)

que = []
D = {}
dcnt = 0
# ポイント上位K個をまず取る
for i in range(K):
    d, t = DT[i]
    dcnt += d
    # 種類tが1つ目
    if t not in D:
        D[t] = d
    # 2つ目以降なら、最大のものはDに残して他はキューへ
    else:
        if D[t] < d:
            D[t], d = d, D[t]
        que.append((d, t))
heapify(que)
tcnt = len(D)

ans = dcnt + tcnt**2
for i in range(K, N):
    d2, t2 = DT[i]
    # 既に持っている種類ならスキップ
    if t2 in D:
        continue
    D[t2] = d2
    tmp = []
    # 全部1種類ずつならもう種類数を増やせないので終了
    if not que:
        break
    d, t = heappop(que)
    # この候補と交換してみる
    dcnt -= d
    dcnt += d2
    tcnt += 1
    ans = max(ans, dcnt + tcnt**2)
print(ans)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
TD=[list(map(int,input().split())) for _ in range(n)]
TD=sorted(TD,key=lambda x:-x[1])
#print(TD)
q=[]
ans=0
v=set()
cnt=0
for t,d in TD[:k]:
    ans+=d
    if t in v:
        q.append(d)
    else:
        v.add(t)

ans+=len(v)**2
cnt=ans
for t,d in TD[k:]:
    if t not in v and len(q)!=0:
        x=q.pop()
        ans+=d+2*len(v)+1-x
        v.add(t)
        cnt=max(ans,cnt)

print(cnt)

Submission status: Accepted

Solution code:
def solve(n, k, t, d):
    que = []
    used = set()
    base_score = 0
    for ti, di in sorted(zip(t, d), key=lambda _:-_[1])[:k]:
        if ti in used:
            que.append(di)
        else:
            used.add(ti)
        base_score += di
    num = len(used)
    best = base_score + num**2

    M = {}
    for ti, di in zip(t, d):
        if not ti in used:
            if not ti in M:
                M[ti] = 0
            M[ti] = max(M[ti], di)

    que = sorted(que)
    for i, p in enumerate(sorted(M.values())[::-1][:len(que)]):
        base_score += p - que[i]
        num += 1
        best = max(best, base_score + num**2)
    return best

n, k = map(int, input().split())
t = [0] * n
d = [0] * n
for i in range(n):
    t[i], d[i] = map(int, input().split())
print(solve(n, k, t, d))

Submission status: Accepted

Solution code:
# D - Various Sushi

import sys

N, K = map(int, sys.stdin.buffer.readline().split())
td = []
for _ in range(N):
    tmp1, tmp2 = map(int, sys.stdin.buffer.readline().split())
    td.append([tmp1-1, tmp2, -1])

td = sorted(td, key = lambda x: x[1], reverse = True)
t_counter = [0]*N
for idx in range(N):
    t_counter[td[idx][0]] += 1
    td[idx][2] = t_counter[td[idx][0]]

t_set = set([td[idx][0] for idx in range(K)])
ans_list = []
ans = sum([td[idx][1] for idx in range(K)]) + len(t_set)**2
ans_list.append(ans)

a = 0
b = 0
while 1:
    while K - 1 - a >= 0:
        if td[K - 1 - a][2] >= 2:
            ans -= td[K - 1 - a][1]
            a += 1
            break
        a += 1
    if K-1-a < 0:
        break
    while K + b <= N-1:
        if td[K + b][0] not in t_set:
            ans -= len(t_set)**2
            t_set.add(td[K + b][0])
            ans += td[K + b][1] + len(t_set)**2
            b += 1
            break
        b += 1
    ans_list.append(ans)
    if K+b == N:
        break

print(max(ans_list))

Submission status: Accepted

Solution code:
def solve():
    from collections import deque, namedtuple
    from operator import attrgetter
    import sys

    input = sys.stdin.readline
    Sushi = namedtuple('Sushi', 'taste kind')

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

    e = []
    for _ in range(n):
        t, d = map(int, input().split())
        e.append(Sushi(taste=d, kind=t - 1))

    e.sort(key=attrgetter('taste'), reverse=True)

    ret = 0
    used = [False] * n  # 選択されている種類
    q = deque()  # 選択から外す候補,[high->low]
    x = 0  # 暫定選択の種類数
    for s in e[:k]:  # 美味しさ降順k個
        if used[s.kind]:
            q.append(s)  # 同じ種類の2個目以降のみ交換候補になる
        else:
            used[s.kind] = True
            x += 1
        ret += s.taste
    ret += pow(x, 2)  # 種類ボーナス

    t = ret
    r = iter(e[k:])  # 交換で入れる候補
    while q:
        rem = q.pop()  # 選択から外す
        t -= rem.taste
        for s in r:
            if used[s.kind]:
                continue
            used[s.kind] = True
            t += s.taste + x * 2 + 1
            x += 1
            ret = max(ret, t)
            break
        else:
            break

    return ret


if __name__ == '__main__':
    print(solve())

# k個持っておいて、下位を交換する

Submission status: Accepted

Solution code:
def solve():
    from collections import namedtuple
    from heapq import heappush, heappop
    from operator import attrgetter
    import sys

    input = sys.stdin.readline
    Sushi = namedtuple('Sushi', 'taste kind')

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

    e = []
    for _ in range(n):
        t, d = map(int, input().split())
        e.append(Sushi(taste=d, kind=t - 1))

    e.sort(key=attrgetter('taste'), reverse=True)

    b = e[:k]  # 美味しさ降順k個

    ret = 0
    used = [False] * n  # 選択されている種類
    h = []  # 選択から外す候補
    x = 0  # 暫定選択の種類数
    for s in b:
        if used[s.kind]:
            heappush(h, s)
        else:
            used[s.kind] = True
            x += 1
        ret += s.taste

    ret += pow(x, 2)  # 種類ボーナス

    t = ret
    r = iter(e[k:])  # 交換で入れる候補
    while h:
        rem = heappop(h)  # 選択から外す
        t -= rem.taste

        for s in r:
            if used[s.kind]:
                continue
            used[s.kind] = True
            t += s.taste + x * 2 + 1
            x += 1
            ret = max(ret, t)
            break
        else:
            break

    return ret


if __name__ == '__main__':
    print(solve())

Submission status: Accepted

Solution code:
# https://atcoder.jp/contests/abc116/submissions/5667980
# 写経

from heapq import heappush, heappop
import sys

input = sys.stdin.readline

n, k = map(int, input().split())
kind_value = [[] for _ in range(n)]
for _ in range(n):
    t, d = map(int, input().split())
    kind_value[t - 1].append(d)

for lis in kind_value:
    lis.sort()

kind_value = sorted(
    (lis for lis in kind_value if lis),
    key=lambda x: x[-1], reverse=True)

ret = 0
curr = 0
kind = 0
remains = []
for lis in kind_value:
    kind += 1
    if kind > k:
        break
    curr += lis[-1]
    for d in lis[:-1]:
        heappush(remains, d)
        curr += d
    while len(remains) > k - kind:
        curr -= heappop(remains)
    ret = max(ret, curr + kind * kind)

print(ret)

Submission status: Accepted

Solution code:
from collections import defaultdict
import heapq
import sys
input = sys.stdin.readline


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

sushi.sort(key=lambda x: x[1])
kinds = defaultdict(int)
eaten = []
ans = 0
for i in range(k):
    neta, point = sushi.pop()
    kinds[neta] += 1
    ans += point
    heapq.heappush(eaten, [point, neta])

num_of_kinds = len(kinds.keys())
ans += num_of_kinds ** 2
candidate = ans

for i in range(k):
    point, neta = heapq.heappop(eaten)
    if kinds[neta] == 1:
        continue
    kinds[neta] -= 1

    while sushi:
        new_neta, new_point = sushi.pop()
        if kinds[new_neta] >= 1:
            continue
        else:
            break
    else:
        break
    kinds[new_neta] += 1
    num_of_kinds += 1
    delta = new_point - point + (2*num_of_kinds - 1)
    candidate += delta
    ans = max(ans, candidate)

print(ans)

Submission status: Accepted

Solution code:
from collections import defaultdict
import heapq


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

sushi.sort(key=lambda x: x[1])
kinds = defaultdict(int)
eaten = []
ans = 0
for i in range(k):
    neta, point = sushi.pop()
    kinds[neta] += 1
    ans += point
    heapq.heappush(eaten, [point, neta])

num_of_kinds = len(kinds.keys())
ans += num_of_kinds ** 2
candidate = ans

for i in range(k):
    point, neta = heapq.heappop(eaten)
    if kinds[neta] == 1:
        continue
    kinds[neta] -= 1

    while sushi:
        new_neta, new_point = sushi.pop()
        if kinds[new_neta] >= 1:
            continue
        else:
            break
    else:
        break
    kinds[new_neta] += 1
    num_of_kinds += 1
    delta = new_point - point + (2*num_of_kinds - 1)
    candidate += delta
    ans = max(ans, candidate)

print(ans)

Submission status: Accepted

Solution code:
import sys
from operator import itemgetter
read = sys.stdin.read

N, K, *td = map(int, read().split())
td = sorted(zip(*[iter(td)] * 2), key=itemgetter(1), reverse=True)

stack = []
taste = 0
neta_group = set()
for t, d in td[:K]:
    taste += d
    if t in neta_group:
        stack.append(d)
    else:
        neta_group.add(t)

variety = len(neta_group)
candidates = [taste + variety ** 2]

if N == K:
    print(candidates[0])
    exit()

rest = list(reversed(td[K::]))

while rest and stack:
    t1, d1 = rest.pop()
    if t1 not in neta_group:
        neta_group.add(t1)
        d2 = stack.pop()
        variety += 1
        taste = taste - d2 + d1

        candidates.append(taste + variety ** 2)

print(max(candidates))

Submission status: Accepted

Solution code:
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
selected0 = set()
selected1 = set()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        if t in selected0:
            heapq.heappush(ls0, (d, t))
        selected0.add(t)
        base += d
    else:
        if t not in selected0 and t not in selected1:
            heapq.heappush(ls1, (-d, t))
            selected1.add(t)

m = len(selected0)
res = base + m**2

for _ in range(k):
    if len(ls0) == 0 or len(ls1) == 0:
        break
    d, _ = heapq.heappop(ls0)
    base -= d
    d, _ = heapq.heappop(ls1)
    base += -d
    m += 1
    res = max(res, base + m**2)
print(res)

Submission status: Accepted

Solution code:
import collections
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
counter0 = collections.Counter()
counter1 = collections.Counter()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        if counter0[t] > 0:
            heapq.heappush(ls0, (d, t))
        counter0[t] += 1
        base += d
    else:
        if t not in counter0 and t not in counter1:
            heapq.heappush(ls1, (-d, t))
            counter1[t] += 1

m = len(counter0)
res = base + m**2

for _ in range(k):
    if len(ls0) == 0 or len(ls1) == 0:
        break
    d, _ = heapq.heappop(ls0)
    base -= d
    d, _ = heapq.heappop(ls1)
    base += -d
    m += 1
    res = max(res, base + m**2)
print(res)

Submission status: Accepted

Solution code:
import collections
import heapq
import operator
import sys

input = sys.stdin.readline

ri = lambda: int(input())
rs = lambda: input().rstrip()
ril = lambda: list(map(int, input().split()))
rsl = lambda: input().rstrip().split()
ris = lambda n: [ri() for _ in range(n)]
rss = lambda n: [rs() for _ in range(n)]
rils = lambda n: [ril() for _ in range(n)]
rsls = lambda n: [rsl() for _ in range(n)]

n, k = ril()
ls = rils(n)
ls.sort(key=operator.itemgetter(1), reverse=True)
counter = collections.Counter()
ls0 = []
ls1 = []
base = 0
for i, (t, d) in enumerate(ls):
    if i < k:
        heapq.heappush(ls0, (d, t))
        counter[t] += 1
        base += d
    else:
        if t not in counter:
            heapq.heappush(ls1, (-d, t))

m = len(counter)
res = base + m**2

for i in range(k):
    ok = False
    while ls0:
        d, t = heapq.heappop(ls0)
        counter[t] -= 1
        if counter[t] > 0:
            base -= d
            ok = True
            break
    if not ok:
        break

    ok = False
    while ls1:
        d, t = heapq.heappop(ls1)
        if t not in counter:
            counter[t] += 1
            base += -d
            ok = True
            break
    if not ok:
        break

    res = max(res, base + (m + i + 1)**2)
print(res)

Submission status: Accepted

Solution code:
from collections import defaultdict

N, K = map(int, input().split())
TD = [tuple(map(int, input().split())) for _ in range(N)]
TD.sort(key=lambda A: A[1], reverse=True)

A = TD[:K][:: -1]
cnt = defaultdict(int)
V = set()
dSum = 0
for t, d in A:
    dSum += d
    cnt[t] += 1
    V.add(t)

i = 0
ans = dSum + len(V)**2
for t, d in TD[K:]:
    if not t in V:
        while i < K and cnt[A[i][0]] <= 1:
            i += 1
        if i == K:
            break

        dSum = dSum - A[i][1] + d
        V.add(t)
        cnt[A[i][0]] -= 1
        i += 1
        ans = max(ans, dSum + len(V)**2)

print(ans)

Submission status: Accepted

Solution code:
n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x: -x[1])

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    type, taste = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        same.append(taste)

rem = []
typesRem = set()
for i in range(nTake, n):
    type, taste = a[i]
    if type not in types and type not in typesRem:
        typesRem.add(type)
        rem.append(taste)

same.sort()
rem.sort(key=lambda x: -x)
ret = totalTaste + len(types) ** 2
curTotal = totalTaste
for i in range(min(len(same), len(rem))):
    curTotal = curTotal - same[i] + rem[i]
    ret = max(ret, curTotal + (len(types) + i + 1) ** 2)
print(ret)

Submission status: Accepted

Solution code:
n, k = map(int, input().split())
td = list(sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda td: -td[1]))
td2 = []
eaten = set()
dup = 0
score = k * k
i = 0
j = 0
while len(eaten) + dup < k:
    if i < n and td[i][0] in eaten:
        td2.append(td[i])
        i += 1
    else:
        diff = (k - dup) * (k - dup) - (k - dup - 1) * (k - dup - 1)
        if i < n and (j >= len(td2) or td[i][1] >= td2[j][1] - diff):
            t, d = td[i]
            score += d
            eaten.add(t)
            i += 1
        else:
            t, d = td2[j]
            score += d - diff
            dup += 1
            j += 1
print(score)

Submission status: Accepted

Solution code:
from heapq import heappop, heappush

n, nTake = map(int, input().split())
a = sorted([tuple(map(int, input().split())) for _ in range(n)], key=lambda x: (-x[1], x[0]))

totalTaste = 0
same = []
types = set()
for i in range(nTake):
    type, taste = a[i]
    totalTaste += taste
    if type not in types:
        types.add(type)
    else:
        heappush(same, taste)

rem = []
typesRem = set()
for i in range(nTake, n):
    type, taste = a[i]
    if type not in types and type not in typesRem:
        typesRem.add(type)
        heappush(rem, -taste)

ret = totalTaste + len(types) ** 2
curTotal = totalTaste
add = 1
while same and rem:
    worst = heappop(same)
    best = heappop(rem)
    curTotal = curTotal - worst - best
    ret = max(ret, curTotal + (len(types) + add) ** 2)
    add += 1
print(ret)

Submission status: Accepted

Solution code:
from collections import*
n,k=map(int,input().split())
sushi=sorted([tuple(map(int,input().split()))for i in range(n)],key=lambda x:x[1],reverse=1)
first=defaultdict(list)
second=defaultdict(list)
menu=set()

for _,v in sushi[:k]:
    menu.add(_)
    first[_].append(v)
second=sushi[k:][::-1]
score=sum(map(sum,first.values()))
change=[]
for _,v in first.items():
    for i in v[1:]:
        change.append((_,i))
change.sort(key=lambda x:x[1],reverse=1)

#print(first,change,second)
ans=[score+len(first)**2]
while change and second:
    _,s=change.pop()
    score-=s
    try:
        while second[-1][0] in menu:
            pp,ss=second.pop()
        pp,ss=second.pop()
        menu.add(pp)
        score+=ss
        ans.append(score+len(menu)**2)
    except:
        break
print(max(ans))

Submission status: Accepted

Solution code:
from heapq import heappush, heappop
N, K = map(int, input().split())
log = [0]*N
pre = []
ls = []
for i in range(N):
  t, d = map(int, input().split())
  t = t-1
  heappush(ls,(-d,t))
v = 0
n = 0
for i in range(K):
  d, t = heappop(ls)
  d = -d
  log[t] += 1
  v += d
  if log[t]>1:
    heappush(pre,(d,t))
  else:
    n += 1
v += n*n
ans = v
while ls and pre:
  d, t = heappop(ls)
  if log[t]>0:
    continue
  log[t] += 1
  d = -d
  pv = v
  x, y = heappop(pre)
  v -= x
  v += d
  v += 2*n+1
  n += 1
  if ans<v:
    ans = v
print(ans)

Submission status: Accepted

Solution code:
import sys
sys.setrecursionlimit(2147483647)
INF=float("inf")
MOD=10**9+7
input=lambda :sys.stdin.readline().rstrip()
def resolve():
    from heapq import heappop,heappush
    n,k=map(int,input().split())
    A=[0]*n
    for i in range(n):
        t,d=map(int,input().split())
        t-=1
        A[i]=(t,d)
    A.sort(lambda x:-x[1]) # d の降順ソート

    # 各種の中で最大かどうかをわける
    B=[0]*n
    first=[1]*n
    for i,a in enumerate(A):
        t,d=a
        if(first[t]):
            first[t]=0
            B[i]=(d,1)
        else:
            B[i]=(d,0)

    # 初めは貪欲に d の大きい方から k 個取る
    S=[]
    for i in range(k):
        d,f=B[i]
        heappush(S,(f,d))
    x=sum(S[i][1] for i in range(k))
    y=sum(S[i][0] for i in range(k))
    ans=x+y*y

    # 残りの n-k について、Sからf=0で最小のものを捨て、B[k:]からf=1の中で最大のものを加える
    B=B[k:]
    for i in range(n-k):
        d,f=B[i]
        if(f==0): continue
        pf,pd=heappop(S)
        if(pf==1): break
        x+=d-pd
        y+=1
        ans=max(ans,x+y*y)
    print(ans)
resolve()

Submission status: Accepted

Solution code:
from operator import itemgetter

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

TD.sort(key=itemgetter(1), reverse=True)
selected = TD[:K]
dlc = sum(map(itemgetter(1), selected))
variety = set(map(itemgetter(0), selected))

used = set()
overlap = [] # ダブっているネタの美味しさのリスト
for t, d in TD[:K]:
    if t not in used:
        used.add(t)
    else:
        overlap.append(d)
answer = dlc + len(variety)**2

for t, d in TD[K:]:
    if not overlap:
        break
    if t in used:
        continue
    ov = overlap.pop()
    dlc = dlc - ov + d
    used.add(t)
    answer = max(answer, dlc + len(used)**2)

print(answer)

Submission status: Accepted

Solution code:
from heapq import heappush, heappop

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

counter = [0] * (N + 1)
h = []
for t, d in td :
    counter[t] += 1
    heappush(h, (d, d, t))

kind = len(set(t for t, d in td))
for _ in range(N - K) :
    while True :
        _, d, t = heappop(h)
        if counter[t] == 1 :
            loss = kind * kind - (kind - 1) * (kind - 1) + d
            if loss <= h[0][0] :
                counter[t] -= 1
                kind -= 1
            else :
                heappush(h, (loss, d, t))
                continue
        else :
            counter[t] -= 1
        
        break
        
print(sum(d for _, d, t in h) + kind * kind)

Submission status: Accepted

Solution code:
# 同じ種類だと、点数が高いやつを選ぶ
# 違う種類を選ぶと、ボーナス点が入るので、それをどうするか
# d  = sum(di) + x**x
# 種類xを固定すると？ O(log(N))でいける？
# ※※ 種類が増えるかどうかを降順で見ながらboolで置き換えるという発想。強い。

from collections import defaultdict

di = defaultdict(list)
n, k = map(int, input().split())

sushi = []
kinds = []
for i in range(n):
    t,d = map(int, input().split())
    sushi.append((d,t))
    kinds.append(t)
sushi.sort(reverse=True)
sushidict = defaultdict(int)
n_kind_all = len(set(kinds))

cnt = 0
import heapq
sushiq = []
# まずは上からK個貪欲に取る
for i in range(k):
    d,t = sushi[i]
    heapq.heappush(sushiq,(d,t)) 
    sushidict[t] += 1
    cnt += d

n_kind = len(sushidict)
res = []
res.append(cnt+n_kind**2)

for i in range(k,n):
    d,t = sushi[i]
    # すでにある種類の場合は大丈夫
    if sushidict[t] > 0 : continue
    # 種類が増える時、小さいものから入れ替えていく
    d_min,t_min = heapq.heappop(sushiq)
    # popして消えるようなものはNG
    while sushidict[t_min] == 1 and sushiq:
        d_min, t_min = heapq.heappop(sushiq)
    # popするものがなくなったら終了
    if not sushiq:
        break
    sushidict[t_min] -= 1
    sushidict[t] += 1
    n_kind += 1
    cnt  = cnt - d_min + d
    res.append(cnt+n_kind**2)
    if n_kind == n_kind_all:
        break

print(max(res))
# print(res)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
TD=[list(map(int,input().split())) for _ in range(n)]
TD=sorted(TD,key=lambda x:-x[1])
#print(TD)
q=[]
ans=0
v=set()
cnt=0
for t,d in TD[:k]:
    ans+=d
    if t in v:
        q.append(d)
    else:
        v.add(t)

ans+=len(v)**2
cnt=ans
for t,d in TD[k:]:
    if t not in v and len(q)!=0:
        x=q.pop()
        ans+=d+2*len(v)+1-x
        v.add(t)
        cnt=max(ans,cnt)

print(cnt)

Submission status: Accepted

Solution code:
n,k=map(int,input().split())
a=[]
for _ in range(n):a.append(tuple(list(map(int,input().split()))[::-1]))
a.sort(reverse=1)
x,y=[],[]
s=set()
for i in range(k):
  d,t= a[i]
  if t in s:
    x.append(d)
  else:
    y.append(d)
    s.add(t)
xx,yy,v=sum(x),sum(y),len(y)
ans=xx+yy+v**2
for i in range(k,n):
    if len(x)==0:break
    d,t=a[i]
    if t in s:continue
    s.add(t)
    xx-=x.pop()
    yy+=d
    v+=1
    ans=max(ans,xx+yy+v**2)
print(ans)

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline
sys.setrecursionlimit(pow(10, 6))
from heapq import heappush, heappop

def main():
    n, k = map(int, input().split())
    td = [tuple(map(int, input().split())) for _ in range(n)]
    td = list(sorted(td, key=lambda x: -x[1]))
    dsum, kind = 0, 0
    num = [0 for _ in range(n)]
    h = []
    for t, d in td[:k]:
        dsum += d
        if num[t-1] > 0:
            heappush(h, (d, t))
        if num[t-1] == 0:
            kind += 1
        num[t-1] += 1
    
    ans = dsum + kind * kind

    for t, d in td[k:]:
        if len(h) == 0:
            break
        if num[t-1] > 0:
            continue
        else:
            _d, _t = heappop(h)
            num[_t-1] -= 1
            num[t-1] += 1
            kind += 1
            dsum += d-_d
            if dsum + kind * kind > ans:
                ans = dsum + kind * kind
    
    print(ans)


if __name__ == '__main__':
    main()

Submission status: Accepted

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

l=[] # 寿司リスト。ただし，「おいしさ」を第1要素にする
m=[] # 寿司タプルをそのまま入れる
s=set()

for i in range(N):
    t,d=map(int,input().split())
    l.append((d,t))
l.sort(reverse=True) 

ans=K**2
t=0 # 既に食べたことのある種類の寿司を更に食べた回数

i=0 # l内を走査する際の添字
j=0 # m内を操作する際の添え字

for _ in range(K): # _番目に食べる寿司について
    while i<N and l[i][1] in s:
        m.append(l[i]) # そのままmに追加する
        i+=1 # lの先頭の要素を削除せず，
    if (i<N and (m and j<len(m)) and l[i][0]>=m[j][0]-(2*K-1-2*t)) or (not m or j>=len(m)):
        ans+=l[i][0]
        s.add(l[i][1])
        i+=1
    else:
        ans+=m[j][0]-(2*K-1-2*t)
        j+=1
        t+=1

print(ans)

Submission status: Accepted

Solution code:
from heapq import heappop, heappush
 
N, K = map(int, input().split())
td = sorted([list(map(int, input().split()))
             for _ in range(N)], key=lambda x: x[1], reverse=True)

q = []
v = set()
count = 0
for t, d in td[:K]:
    count += d
    if t in v:
        heappush(q, d)
    else:
        v.add(t)
count += len(v) ** 2
ans = count
for t, d in td[K:]:
    if t not in v and len(q) != 0:
        x = heappop(q)
        count += d + 2 * len(v) + 1 - x
        v.add(t)
        ans = max(ans, count)
        
print(ans)

Submission status: Accepted

Solution code:
import sys
n,k=map(int,input().split())
data=[]
for i in range(n):
    t,d=map(int,input().split())
    data.append([d,t])
data.sort(reverse=True)
flag=[0]*(n+1)
eat=[]
lst=[]
count=0
for u in data:
    d,t=u
    if flag[t]==0:
        eat.append(u)
        flag[t]=1
        count+=1
        if count==k:
            break
    else:
        lst.append(u)
else:
    lst.sort(reverse=True)
    sum=0
    for i in range(n-count):
        if i<k-count:
            sum+=lst[i][0]
        else:
            if lst[i][0]>eat[count-1][0]+2*count-1:
                eat[count-1]=lst[i]
                count-=1
            else:
                break
    for u in eat:
        sum+=u[0]
    sum+=count**2
    print(sum)
    sys.exit()
lst.sort(reverse=True)
for u in lst:
    d=u[0]
    if d>eat[count-1][0]+2*count-1:
        eat[count-1]=u
        count-=1
    else:
        break
sum=0
for u in eat:
    sum+=u[0]
sum+=count**2
print(sum)

Submission status: Accepted

Solution code:
import heapq

INF=10**12

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

for i in range(N):
    td[i][0], td[i][1] = map(int,input().split())
    td[i][0] -= 1
    s[td[i][0]].append(td[i][1])

for i in range(N):
    if s[i] == []:
        s[i].append(-INF)
    s[i].sort(key=lambda x:-x)

s.sort(key=lambda x:-x[0])

tmp = 0
q = []

for i in range(K):
    tmp += s[i][0]
    for j in range(1,len(s[i])):
        heapq.heappush(q,-s[i][j])
for i in range(K,N):
    for j in range(len(s[i])):
        heapq.heappush(q,-s[i][j])

res = tmp + K * K

for i in range(1,K)[::-1]:
    v = s[i][0]
    if q == []: break
    w = - q[0]
    if v < w:
        tmp = tmp + w - v
        heapq.heappop(q)
        heapq.heappush(q,-v)
    res = max(res, tmp + i * i)        

print(res)

Submission status: Accepted

Solution code:
from heapq import heappush, heapreplace

def main():
    INF = 10 ** 9 + 5

    N, K, *TD = map(int, open(0).read().split())

    E = [[] for _ in range(N)]
    for t, d in zip(*[iter(TD)] * 2):
        E[t - 1].append(d)

    for e in E:
        if e:
            e.sort(reverse=True)
        else:
            e.append(-INF)
    E.sort(key=lambda x: -x[0])

    cur = 0
    Q = []
    for ei in E[:K]:
        cur += ei[0]
        for eij in ei[1:]:
            heappush(Q, -eij)

    for ei in E[K:]:
        for eij in ei:
            heappush(Q, -eij)

    res = cur + K * K
    for x in reversed(range(1, K)):
        v, w = E[x][0], -Q[0]
        if v < w:
            heapreplace(Q, -v)
            cur += w - v

        tmp = cur + x * x
        if res < tmp:
            res = tmp

    print(res)

if __name__ == '__main__':
    main()