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: Wrong Answer

Solution code:
N,K = (int(x) for x in input().split())
td_arr = []
for _ in range(N):
    td = tuple(int(x) for x in input().split())
    td_arr.append(td)

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

n_dm_set = set()
small_arr = []
d_dm = 0
for t,d in td_arr[:K]:
    if t in n_dm_set:
        small_arr.append(d)
    n_dm_set.add(t)
    d_dm += d

candidate_dict = {}
candidate_dict[len(n_dm_set)] = d_dm

prev_t = -1
other_arr = []
for t,d in sorted(td_arr, key=lambda x:x[0]):
    if t == prev_t:
        continue
    if t in n_dm_set:
        continue
    other_arr.append((t,d))

d_x = d_dm
nt_x = len(n_dm_set)
for t,d in other_arr:
    if len(small_arr) == 0:
        break
    d_x += d
    d_x -= small_arr[-1]
    small_arr = small_arr[:-1]
    nt_x += 1
    candidate_dict[nt_x] = d_x

ans = 0
for k,v in candidate_dict.items():
    ans = max(ans, k*k+v)

print(ans)

Submission status: Wrong Answer

Solution code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
delicious=0
kind=set()
stack=[]
for t,d in TD[:K]:
    delicious+=d
    if t in kind:stack.append(t)
    else:kind.add(d)
ans1=delicious+len(kind)**2
ans2=ans1
for t,d in TD[K:]:
    if not stack:break
    if not t in kind:
        ans2+=d-stack.pop()+2*len(kind)+1
        kind.add(t)
        ans2=max(ans1,ans2)
print(ans2)

Submission status: Wrong Answer

Solution code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
delicious=0
kind=set()
stack=[]
for t,d in TD[:K]:
    delicious+=d
    if not t in kind:kind.add(t)
    else:stack.append(d)
ans1=delicious+len(kind)**2
ans2=ans1
for t,d in TD[K:]:
    if not stack:break
    if not t in kind:
        ans2+=d-stack.pop()+2*len(kind)+1
        kind.add(t)
        ans2=max(ans1,ans2)
print(ans2)

Submission status: Wrong Answer

Solution code:
N,K=map(int,input().split())
TD=[[int(j) for j in input().split()] for i in range(N)]
TD=sorted(TD,key=lambda x:-x[1])
delicious=0
kind=set([])
surplus=[]
for t,d in TD[:K]:
    delicious+=d
    if not t in kind:kind.add(t)
    else:surplus.append(d)
ans1=delicious+len(kind)**2
ans2=ans1
for t,d in TD[K:]:
    if not surplus:break
    if not t in kind:
        ans2+=2*len(kind)+1
        ans2+=d
        ans2-=surplus.pop()
        kind.add(t)
print(max(ans1,ans2))

Submission status: Wrong Answer

Solution code:
N, K = [int(_) for _ in input().split()]
TD = [[int(_) for _ in input().split()] for i in range(N)]
 
# print(TD)
 
xs = sorted(TD, key=lambda x:-x[1])
# reverseじゃなくてもいい
 
#print(xs)
 
baset = 0
ts = set([])
ds = []
 
for t, d in xs[:K]:
    baset += d
    if not t in ts:
        ts.add(t)
    else:
        ds.append(d)
 
result = baset + len(ts) ** 2
r = result
 
#print(r, ts, ds)
 
for t, d in xs[K:]:
    if not ds:# リストが空かどうか
        print("空")
        break
    if not t in ts:
        r += 2 * len(ts) + 1
        r += d - ds.pop()
        ts.add(t)
        result = max(result, r)
        #print(r, ts, ds)
 
#print(r, ts, ds)
 
print(result)

Submission status: Wrong Answer

Solution code:
from heapq import heappush, heappop

N, K = map(int, input().split(' '))
h = []
for _ in range(N):
    t, d = map(int, input().split(' '))
    heappush(h, (-d, t))

dsum = 0
used = {}
q = []
for _ in range(K):
    d, t = heappop(h)
    if t not in used:
        used[t] = 1
    else:
        heappush(q, (-d, t))
    dsum += -d
res = len(used)**2 + dsum
if len(used) == K:
    print(res)

while h and q:
    d, t = heappop(h)
    if t not in used:
        used[t] = 1
        m, _ = heappop(q)
        dsum = dsum - m - d
        res = max(res, len(used)**2 + dsum)
print(res)

Submission status: Wrong Answer

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], reverse = True)
eat_tops = []
eat_others = []
eat_types = set([])
for t, d in td[:k]:
    if not t in eat_types:
        eat_types.add(t)
        eat_tops.append((t, d))
    else:
        eat_others.append((t, d))
ans = sum([d for _, d in td[:k]]) + len(eat_types)**2
answers = [ans]

for t, d in td[k:]:
    if t in eat_types:
        pass
    else:
        eat_types.add(t)
        eat_tops.append((t, d))
        eat_others = eat_others[:-1]
        ans = sum([d for _, d in eat_tops+eat_others]) + len(eat_types)**2
        answers.append(ans)
print(max(answers))

Submission status: Wrong Answer

Solution code:
N,K = map(int,input().split())
lis = []
val = 0
ans = 0
s = set()
dub = []

for _ in range(N):
    lis.append(tuple(map(int,input().split())))

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

for t,v in lis[:K]:

    if t in s:
        dub.append(v)
    s.add(t)
    val += v

ans = val + len(s) ** 2

for t,v in lis[K:]:
    if len(dub) == 0:
        break
    if t in s:
        continue

    s.add(t)
    r = dub.pop()
    val += v - r
    ans = max(ans,val+len(s) ** 2)

    print(ans)

Submission status: Wrong Answer

Solution code:
import itertools
import heapq
N, K = [int(x) for x in input().split()]
orig_td = [tuple([int(x) for x in input().split()]) for _ in range(N)]
td = sorted(orig_td.copy(), key=lambda x: x[1], reverse=True)
INF = 10**9
# print(td)


def score(a):
    kiso = sum([x[1] for x in a])
    bonus = len(set([x[0] for x in a]))**2
    return kiso + bonus


first = []
second = []
kind = set()
for i in range(K):
    if td[i][0] not in kind:
        first.append(td[i])
        kind.add(td[i][0])
    else:
        second.append(td[i])

best_ans = sum([x[1] for x in first + second]) + len(first)**2

for k in range(K, N):
    if len(second) == 0:
        break
    if td[k] in kind:
        continue

    a = second.pop(-1)
    kind.add(td[k][0])
    first.append(td[k])
    ans = sum([x[1] for x in first + second]) + len(first)**2
    # print(first, second, ans)
    if best_ans < ans:
        best_ans = ans
print(best_ans)

Submission status: Wrong Answer

Solution code:
import itertools
import heapq
N, K = [int(x) for x in input().split()]
orig_td = [tuple([int(x) for x in input().split()]) for _ in range(N)]
td = sorted(orig_td.copy(), key=lambda x: x[1])
INF = 10**9
# print(td)


def score(a):
    kiso = sum([x[1] for x in a])
    bonus = len(set([x[0] for x in a]))**2
    return kiso + bonus


kutta = td[-K:]                 # 貪欲にK個
td = td[:-K]

kutta_kind = set([x[0] for x in kutta])        # 食った種類
kutta_kazu = [0 for _ in range(N+1)]             # 種類ごと食った数
for neta in kutta:
    kutta_kazu[neta[0]] += 1
td = [x for x in td if not x[0] in kutta_kind]  # kutta_kindにあるなら、不要
td_kind = set([x[0] for x in td])               # 食ってない種類

ans = [score(kutta)]
for _ in range(N):
    a = [k for k in kutta if kutta_kazu[k[0]] > 1]
    if len(a) == 0:
        break
    a = min(a, key=lambda x: x[0])
    kutta.remove(a)
    kutta_kazu[a[0]] -= 1

    a = [k for k in td if kutta_kazu[k[0]] == 0]
    if len(a) == 0:
        break
    a = max(a, key=lambda x: x[1])
    kutta.append(a)
    td.remove(a)
    kutta_kazu[a[0]] += 1

    ans.append(score(kutta))
print(max(ans))

# print("全列挙", max([score(x) for x in itertools.permutations(orig_td, K)]))

Submission status: Wrong Answer

Solution code:
N, K = [int(x) for x in input().split()]
td = [tuple([int(x) for x in input().split()]) for _ in range(N)]
INF = 10**9
print(td)


def score(a):
    kiso = sum([x[1] for x in a])
    # print(a)
    bonus = len(set([x[0] for x in a]))**2
    return kiso + bonus


total = score(td)
for i in range(N, K, -1):
    max_score = -INF
    useless_neta = 0
    for neta in td:
        te = td.copy()
        te.remove(neta)
        kari = score(te)
        if kari > max_score:
            max_score = kari
            useless_neta = neta
    td.remove(useless_neta)
    print(td)
print(score(td))

Submission status: Wrong Answer

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

T = 0
neta = set()
doublecounts = []
ans = []

for i in range(K):
    T += TD[i][1]    
    neta.add(TD[i][0])
    if (TD[i][0] in neta):       
        doublecounts.append(TD[i][1])    
            
ans.append(T+len(neta)**2)

for i in range(K,N):
    if(len(doublecounts)==0):
        break
    if not(TD[i][0] in neta):
        T -= doublecounts.pop(-1)
        T += TD[i][1]
        neta.add(TD[i][0])   
        ans.append(T+len(neta)**2)
print(max(ans))

Submission status: Wrong Answer

Solution code:
class Person:
    def __init__(self):
        self.had_set = set([])
        self.satisfied_t = 0
        self.satisfied_d = 0

    def have(self,t,d):
        self.satisfied_d += d
        self.had_set = self.had_set | {t}
        self.satisfied_t = len(self.had_set)**2
        return
    
    def next_satisfy(self):
        before = self.satisfied_t
        after = (len(self.had_set)+1)**2
        return after - before

N, K = map(int, input().split())
list_t, list_d = [],[]
for i in range(N):
    s1,s2 = map(int, input().split())
    list_t.append(s1)
    list_d.append(s2)

I = Person()

for i in range(K):
    remaining_t = set(list_t)-I.had_set
    idx = list_d.index(max(list_d))
    if len(remaining_t)>0:
        if list_d[idx] > I.next_satisfy():
            I.have(list_t.pop(idx),list_d.pop(idx))
        else:
            idxs_r = [i for i,t in enumerate(list_t) if t in remaining_t]
            list_r_t = [t for i,t in enumerate(list_d) if i in idxs_r]
            list_r_d = [d for i,d in enumerate(list_d) if i in idxs_r]
            idx_of_idxs = list_r_d.index(max(list_r_d))
            idx = idxs_r[idx_of_idxs]
            I.have(list_t.pop(idx),list_d.pop(idx))
    else:
        I.have(list_t.pop(idx),list_d.pop(idx))
print(I.satisfied_t + I.satisfied_d)

Submission status: Wrong Answer

Solution code:
from collections import deque


def main():
    N, K = map(int, input().split())
    sushis = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        sushis[i] = (t, d)
    sushis.sort(key=lambda a: (-a[1], a[0]))
    tmp = 0
    types = {}
    q = deque()
    stack = deque()
    for s in sushis:
        t, d = s
        if t in types:
            q.append((t, d))
            continue
        types[t] = 1
        tmp += d
        stack.append((t, d))
        if len(types) >= K:
            break
    cnt = len(types)
    ans = tmp + cnt**2
    while len(q) and len(stack):
        if cnt >= K:
            t, d = stack.pop()
            tmp -= d
            types[t] -= 1
            if types[t] == 0:
                del types[t]
        t, d = q.popleft()
        if t not in types:
            types[t] = 0
        types[t] += 1
        tmp += d
        if tmp + len(types)**2 > ans:
            ans = tmp + len(types)**2
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
from collections import deque


def main():
    N, K = map(int, input().split())
    sushis = [None] * N
    for i in range(N):
        t, d = map(int, input().split())
        sushis[i] = (t, d)
    sushis.sort(key=lambda a: (-a[1], a[0]))
    tmp = 0
    types = {}
    q = deque()
    stack = deque()
    for s in sushis:
        t, d = s
        if t in types:
            q.append((t, d))
            continue
        types[t] = 1
        tmp += d
        stack.append((t, d))
        if len(types) >= K:
            break
    cnt = len(types)
    ans = tmp + cnt**2
    while len(q):
        if cnt >= K:
            t, d = stack.pop()
            tmp -= d
            types[t] -= 1
            if types[t] == 0:
                del types[t]
        t, d = q.popleft()
        if t not in types:
            types[t] = 0
        types[t] += 1
        tmp += d
        if tmp + len(types)**2 > ans:
            ans = tmp + len(types)**2
        else:
            break
    print(ans)


if __name__ == "__main__":
    main()

Submission status: Wrong Answer

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, vari


selected = sushis[:K]
unSelected = sushis[K:]

p, vari = calcPoint(selected)
cand = [p]

va = len(vari)
v = [0]*(N+1)
for i, _ in selected:
    v[i] += 1

can_del = []
for i, d in reversed(selected):
    if v[i] > 2:
        v[i] -= 1
        can_del.append(d)

can_in = []
for i, d in unSelected:
    if not i in vari:
        vari.add(i)
        can_in.append(d)

for od, nd in zip(can_del, can_in):
    nextp = p-od+nd-(va)**2+(va+1)**2
    p = nextp
    cand.append(nextp)
    va += 1
print(max(cand))

Submission status: Wrong Answer

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, vari


selected = sushis[:K]
unSelected = sushis[K:]

p, vari = calcPoint(selected)
cand = [p]

va = len(vari)
v = [0]*(N+1)
for i, _ in selected:
    v[i] += 1

can_del = []
for i, d in reversed(selected):
    if v[i] > 2:
        v[i] -= 1
        can_del.append(d)

can_in = []
for i, d in unSelected:
    if not i in vari:
        vari.add(i)
        can_in.append(d)

for od, nd in zip(can_del, can_in):
    nextp = p-od+nd-(va)**2+(va+1)**2
    cand.append(nextp)
    va += 1
print(max(cand))

Submission status: Wrong Answer

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p


def select_sushi(sushis):
    """
    sushisは美味しさ順にソート済み。種類が最大になるように上から取っていく
    """
    v = set()
    ret_sushis = []
    un_select = []
    for i, d in sushis:
        if not i in v:
            ret_sushis.append([i, d])
            v.add(i)
        else:
            un_select.append([i, d])
    return ret_sushis, un_select, len(v)


select, unSelect, v = select_sushi(sushis)
r = K-len(select)
for i in range(r):
    select.append(unSelect[i])
unSelect = unSelect[r:]
p = calcPoint(select)

vari = [0]*(N+1)
for i, _ in select:
    vari[i] += 1

for i in range(min(len(unSelect), K)):
    si, sd = select[K-1-i]
    ui, ud = unSelect[i]
    update = False
    if vari[si] > 1:
        next_p = p - sd + ud
        if p < next_p:
            vari[si] -= 1
            vari[ui] += 1
            p = next_p
            update = True
    else:
        next_p = p - sd + ud - (v)**2+(v-1)**2
        if next_p > p:
            vari[si] -= 1
            vari[ui] += 1
            p = next_p
            update = True
            v -= 1
    if not update:
        break
print(p)

Submission status: Wrong Answer

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, len(vari)


selected = sushis[:K]
rest_shushi = sushis[K+1:]
p, nv = calcPoint(selected)

vari = [0]*(N+1)
for i, v in selected:
    vari[i] += 1

for i, v in rest_shushi:
    if vari[i] >= 1:
        continue
    else:
        d = []
        for a, b in reversed(selected):
            if vari[a] > 1:
                next_p = p-b+v + (nv+1)**2 - nv**2
                if next_p > p or True:
                    vari[a] -= 1
                    vari[i] += 1
                    nv += 1
                    p = next_p
print(p)

Submission status: Wrong Answer

Solution code:
# python template for atcoder1
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline

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

sushis = sorted(sushis, key=lambda x: (x[1], x[0]), reverse=True)

p = 0


def calcPoint(selected_sushis):
    p = 0
    vari = set()
    for i, v in selected_sushis:
        vari.add(i)
        p += v
    p += len(vari)**2
    return p, len(vari)


selected = sushis[:K]
rest_shushi = sushis[K+1:]
p, nv = calcPoint(selected)

vari = [0]*(N+1)
for i, v in selected:
    vari[i] += 1

for i, v in rest_shushi:
    if vari[i] >= 1:
        continue
    else:
        d = []
        for a, b in reversed(selected):
            if vari[a] > 1:
                next_p = p-b+v + (nv+1)**2 - nv**2
                if next_p > p:
                    vari[a] -= 1
                    vari[i] += 1
                    nv += 1
                    p = next_p
print(p)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())


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

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

#-----------------------------------------
best_list=[]
not_best_list=[]

m=-1
for i in li:
    if m==i[0]:
        not_best_list.append(i)
    else:
        best_list.append(i)
        m=i[0]
        
not_best_list.sort(reverse=True,key=lambda x:x[1])
best_list.sort(reverse=True,key=lambda x:x[1])


s2=min(find_num_neta(list(range(n)),li),k)
s1=find_num_neta(list(range(k)),li)
 
 
score=-1
for i in range(s1,s2+1):
    oishisa=0
    for j in range(i):
        oishisa+=best_list[j][1]
    for j in range(k-i):
        oishisa+=not_best_list[j][1]
    
    new_score=i**2+oishisa
    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Wrong Answer

Solution code:
n, k = [int(x) for x in input().split()]
tops = {}
rest = []
for _ in range(n):
  t, s = [int(x) for x in input().split()]
  if t not in tops:
    tops[t] = s
  elif tops[t] < s: 
    rest.append(tops[t])
    tops[t] = s
  else:
    rest.append(s)
    
tops = [tops[k] for k in tops]
tops.sort(reverse=True)
for i in range(1, len(tops)):
  tops[i] += tops[i-1]
rest.sort(reverse=True)
for i in range(1, len(rest)):
  rest[i] += rest[i-1]

if len(tops) < k:
  i_rest = k - len(tops) - 1
  base = tops[-1] + rest[i_rest]
  max_score = base + len(tops)**2
else:
  i_rest = -1
  base = tops[k-1]
  max_score = base + k**2

i_rest += 1

for i in range(i_rest, min([len(rest), k])):
  base_old = base
  base = tops[k-i-2] + rest[i]
  if (base <= base_old):
    break
  score = base + (k-i-1)**2
  if max_score < score:
    max_score = score

print(max_score)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())
 
 
li.sort(reverse=True,key=lambda x:x[1])
 
def calc_score(li,choice):
    score=0
    k=[]
    for i in choice:
        score+=li[i][1]
        k.append(li[i][0])        
    score+=len(set(k))**2
    return score

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

first_choice=list(range(k))
score=calc_score(li,first_choice)

li_choice_set=set([li[i][0] for i in first_choice])



for i in range(find_num_neta(first_choice,li)+1,min(find_num_neta(list(range(n)),li),k)+1):
  
    new_choice=first_choice.copy()
    change=i-find_num_neta(first_choice,li)
    
    can_append=change
    must_delete=change
    

    
    
    for j in range(k,n):
        if li[j][0] not in li_choice_set:
            new_choice.append(j)
            can_append+=-1
            if can_append==0:
                break
            
            
    for j in range(k-1,-1,-1):

        if li[j][0] in [li[p][0] for p in range(j)]:
            new_choice.remove(j)
            must_delete+=-1
            if must_delete==0:
                break
        
    
    new_score=calc_score(li,new_choice)

    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Wrong Answer

Solution code:
# -*- coding: utf-8 -*-
"""
Created on Fri Feb 22 08:12:39 2019

@author: matsuda
"""

n,k=3,0
li=[[1,9],[3,9],[2,6]]

"""
n,k=7,4
li=[[1,1],
[2,1],
[3,1],
[4,6],
[4,5],
[4,5],
[4,5]]

"""

n,k=6,5
li=[[5,1000000000],
[2,990000000],
[3,980000000],
[6,970000000],
[6,960000000],
[4,950000000]]

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

def calc_score(li,choice):
    score=0
    k=[]
    for i in choice:
        score+=li[i][1]
        k.append(li[i][0])        
    score+=len(set(k))**2
    return score

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

first_choice=list(range(k))
score=calc_score(li,first_choice)

li_choice_set=set([li[i][0] for i in first_choice])



for i in range(find_num_neta(first_choice,li)+1,min(find_num_neta(list(range(n)),li),k)+1):
  
    new_choice=first_choice.copy()
    change=i-find_num_neta(first_choice,li)
    
    can_append=change
    must_delete=change
    

    
    
    for j in range(k,n):
        if li[j][0] not in li_choice_set:
            new_choice.append(j)
            can_append+=-1
            if can_append==0:
                break
            
            
    for j in range(k-1,-1,-1):

        if li[j][0] in [li[p][0] for p in range(j)]:
            new_choice.remove(j)
            must_delete+=-1
            if must_delete==0:
                break
        
    
    new_score=calc_score(li,new_choice)

    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())
 

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

def calc_score(li,choice):
    score=0
    k=[]
    for i in choice:
        score+=li[i][1]
        k.append(li[i][0])        
    score+=len(set(k))**2
    return score

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

first_choice=list(range(k))
score=calc_score(li,first_choice)

li_choice_set=set([li[i][0] for i in first_choice])



for i in range(find_num_neta(first_choice,li)+1,min(find_num_neta(list(range(n)),li)+1,k)):
    new_choice=first_choice.copy()
    change=i-find_num_neta(first_choice,li)
    
    can_append=change
    must_delete=change
    

    
    
    for j in range(k,n):
        if li[j][0] not in li_choice_set:
            new_choice.append(j)
            can_append+=-1
            if can_append==0:
                break
            
            
    for j in range(k-1,-1,-1):

        if li[j][0] in [li[p][0] for p in range(j)]:
            new_choice.remove(j)
            must_delete+=-1
            if can_append==0:
                break
        
    
    new_score=calc_score(li,new_choice)

    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Wrong Answer

Solution code:
n,k=map(int,input().split())
li=[[0,0] for i in range(n)]
for i in range(n):
    li[i][0],li[i][1]=map(int,input().split())

def calc_score(li,choice):
    score=0
    k=[]
    for i in choice:
        score+=li[i][1]
        k.append(li[i][0])        
    score+=len(set(k))**2
    return score

def find_num_neta(choice,li):
    k=[]
    for i in choice:
        k.append(li[i][0])
    return len(set(k))

first_choice=list(range(k))
score=calc_score(li,first_choice)
li_choice_set=set([li[i][0] for i in first_choice])



for i in range(find_num_neta(first_choice,li)+1,find_num_neta(list(range(n)),li)+1):
    new_choice=first_choice.copy()
    change=i-find_num_neta(first_choice,li)
    
    can_append=change
    must_delete=change
    

    
    
    for j in range(k,n):
        if li[j][0] not in li_choice_set:
            new_choice.append(j)
            can_append+=-1
            if can_append==0:
                break
            
            
    for j in range(k-1,-1,-1):

        if li[j][0] in [li[p][0] for p in range(j)]:
            new_choice.remove(j)
            must_delete+=-1
            if can_append==0:
                break
        
    
    new_score=calc_score(li,new_choice)

    
    if score<new_score:
        score=new_score
        
    
print(score)

Submission status: Wrong Answer

Solution code:
#-*- coding: utf-8 -*-
"""
oj dl https://atcoder.jp/contests/abc116/tasks/abc116_d -d test-d
oj test -d test-d -c "python abc116d.py"
oj test -d test-d -c "python abc116d.py" test-d/sample-3.in
"""

from collections import defaultdict
import sys
import time


def deco(func):
    def wrapper(*args, **kwargs):
        s = time.time()
        #print('--start--')
        ret = func(*args, **kwargs)
        logging.debug('--end in %f --' % (time.time() - s))
        return ret

    return wrapper


import itertools
#@deco
def sol_naive(n, k, S):
    maxi = 0
    for comb in itertools.combinations(S, k):
        neta = defaultdict(int)
        ttl = 0
        for t, d in comb:
            ttl += d
            neta[t] += 1
        ttl += len(neta)**2
        maxi = max(maxi, ttl)
    return maxi


def taste(S):
    return sum([d for t, d  in S])
def bonus(S):
    q = len(list(set([t for t, d  in S])))
    return q * q
def sati(S):
    return taste(S) + bonus(S)

def pick_harmless(Sa):
    logging.debug("type(Sa):%s" % type(Sa))
    cache = defaultdict(list)
    for i, (t, d) in enumerate(Sa):
        # most smallest and len(t) > 1:
        cache[t].append(d)
    for k, v in cache.items():
        if len(v) >= 2:
            non_effect_smallest = (k, min(v))
            logging.debug("before Sa:%s" % Sa)
            Sa.remove(non_effect_smallest)
            logging.debug("after  Sa:%s" % Sa)
            return Sa
    return Sa

def add_more_kind(Sa, S):
    Scopied = S.copy()
    for s in Sa:
        Scopied.remove(s)
    n = len(Sa)
    logging.debug("Scopied:%s" % Scopied)
    ts = list(set([ t for t, d in Sa]))
    Ssorted = sorted(Scopied, reverse=True, key=lambda x: x[1])
    for t, d in Ssorted:
        if t not in ts:
            logging.debug("appending maxium with different t: %s, %s" % (t, d))
            Sa.append((t, d))
            logging.debug("after  Sa:%s" % Sa)
            break
    return Sa

def sol(n, k, S):
    logging.debug("S:%s" % S)
    S.sort(key=lambda x: x[1], reverse=True)
    Sa = S[0:k]
    x = len(set([t for (t, d) in Sa]))
    maxi = sati(Sa)
    for i in range(x, k+1):
        # choose min Sushi not reducing the x
        sz_before = len(Sa)
        Sa = pick_harmless(Sa)
        if len(Sa) == sz_before: continue
        # choose max Sushi adding x+1
        Sa = add_more_kind(Sa, S)
        # update maxi
        tmp = sati(Sa)
        logging.debug("tmp:%s, current maxi:%s" % (tmp, maxi))
        maxi = max(maxi, tmp)
    return maxi


import logging

#logging.basicConfig(level=logging.DEBUG, format="%(message)s")
logging.basicConfig(level=logging.ERROR, format="%(message)s")
do_submit = True
#do_submit = False
def input_parse(input_str):
    lines = [x.strip() for x in input_str.split("\n") if x.strip()]
    parsed_lines = [list(map(int, line.split())) for line in lines]
    n, k = parsed_lines[0]
    S = []
    for i in range(n):
        t, d = parsed_lines[i+1]
        S.append((t, d))
    return (n, k, S)


if not do_submit:
    n, k, S = input_parse("""
    5 3
    1 9
    1 7
    2 6
    2 5
    3 1
    """)
    print(sol(n, k, S))
else:
    n, k = list(map(int, input().split()))
    S = []
    for i in range(n):
        t, d = list(map(int, input().split()))
        S.append((t, d))
    print(sol(n, k, S))