Problem p03148 (python)

Fixed submission:
N,K = map(int,input().split())
td = [[int(i) for i in input().split()] for _ in range(N)]
td.sort(key=lambda x:x[1],reverse=True)
ans = 0
ls = set()
dl = []
for i in range(K):
    ans += td[i][1]
    if not td[i][0] in ls:
        ls.add(td[i][0])
    else:
        dl.append(td[i][1])
ans = ans + len(ls)**2
res = ans
for i in range(K,N):
    if not dl:
        break
    if td[i][0] in ls:
        continue
    ans += len(ls)*2+1
    ans += td[i][1]
    ans -= dl[-1]
    del dl[-1]
    res = max(res,ans)
    ls.add(td[i][0])
print(res)

Buggy submission (verdict: Runtime Error):
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)
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)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)
 
L = [list(map(int, input().split())) for i in range(N)]
L.sort(key=lambda x:x[1], reverse=True)
ans = []
etc = []
dic = {}
ind = {}
for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
    if len(dic.keys()) > K:
        break
for t in dic.keys():
    dic[t].sort(reverse=True)
    for i,n in enumerate(dic[t]):
      if i == 0:
          ans.append((n, 1))
      else:
          etc.append((n, 0))
ans_all = calc_ans(ans)
ans.sort(reverse=True)
etc.sort(reverse=True)
for n in etc:
    if ans[-1][1] == 0 and ans[-1][0] >= n[0]:
        break
    ans[-1] = n
    tmp = calc_ans(ans)
    if ans_all < tmp:
        ans_all = tmp
    ans.sort(reverse=True)
print(ans_all)

Buggy submission (verdict: Runtime Error):
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)
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)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

L = [list(map(int, input().split())) for i in range(N)]
L.sort(key=lambda x:x[1], reverse=True)
ans = []
etc = []
dic = {}
ind = {}
for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
    if len(dic.keys()) > K:
        break
for t in dic.keys():
    dic[t].sort(reverse=True)
    for i,n in enumerate(dic[t]):
      if i == 0:
          ans.append((n, 1))
      else:
          etc.append((n, 0))
ans_all = calc_ans(ans)
ans.sort(reverse=True)
etc.sort(reverse=True)
for n in etc:
    #if ans[-1][1] == 0 and ans[-1][0] >= n[0]:
    #    break
    ans[-1] = n
    tmp = calc_ans(ans)
    if ans_all < tmp:
        ans_all = tmp
    ans.sort(reverse=True)
print(ans_all)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

L = [list(map(int, input().split())) for i in range(N)]
L.sort(key=lambda x:x[1], reverse=True)
ans = []
etc = []
dic = {}
ind = {}
for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
    if len(dic.keys()) > K:
        break
for t in dic.keys():
    dic[t].sort(reverse=True)
    for i,n in enumerate(dic[t]):
      if i == 0:
          ans.append((n, 1))
      else:
          etc.append((n, 0))
ans_all = calc_ans(ans)
ans.sort(reverse=True)
etc.sort(reverse=True)
for n in etc:
    if ans[-1][1] == 0 and ans[-1][0] >= n[0]:
        break
    ans[-1] = n
    tmp = calc_ans(ans)
    if ans_all < tmp:
        ans_all = tmp
    ans.sort(reverse=True)
print(ans_all)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

L = [list(map(int, input().split())) for i in range(N)]
L.sort(key=lambda x:x[1], reverse=True)
ans = []
etc = []
dic = {}
for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
    if len(dic.keys()) > K:
        break
for t in dic.keys():
    dic[t].sort(reverse=True)
    for i,n in enumerate(dic[t]):
      if i == 0:
          ans.append((n, 1))
      else:
          etc.append((n, 0))
ans_all = calc_ans(ans)
ans.sort(reverse=True)
for n in etc:
    ans[-1] = n
    tmp = calc_ans(ans)
    if ans_all < tmp:
        ans_all = tmp
    ans.sort(reverse=True)
print(ans_all)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

L = [list(map(int, input().split())) for i in range(N)]
L.sort(key=lambda x:x[1], reverse=True)
for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
    if len(dic.keys()) > K:
        break
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
ans_res = calc_ans(ans)
all_ans = ans_res
for i, n in enumerate(etc):
    ans.sort(reverse=True)
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if all_ans < tmp_res:
        all_ans = tmp_res
    ans_res = tmp_res
print(all_ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
ans_res = calc_ans(ans)
all_ans = ans_res
for i, n in enumerate(etc):
    ans.sort(reverse=True)
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if all_ans < tmp_res:
        all_ans = tmp_res
    ans_res = tmp_res
print(all_ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
ans_res = calc_ans(ans)
all_ans = ans_res
for i, n in enumerate(etc):
    i+=1
    ans[-1*i] = n
    tmp_res = calc_ans(ans)
    if all_ans < tmp_res:
        all_ans = tmp_res
    ans_res = tmp_res
print(all_ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
ans_res = calc_ans(ans)
all_ans = ans_res
for n in etc:
    ans.sort(reverse=True)
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if all_ans < tmp_res:
        all_ans = tmp_res
    ans_res = tmp_res
print(all_ans)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
ans_res = calc_ans(ans)
for n in etc:
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans_res = tmp_res
print(ans_res)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
ans_res = calc_ans(ans)
for n in etc:
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans_res = tmp_res
print(ans_res)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
ans_res = calc_ans(ans)
for n in etc:
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans_res = tmp_res
        ans.sorted(reverse=True)
print(ans_res)

Buggy submission (verdict: Runtime Error):
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
ans_res = calc_ans(ans)
for n in etc:
    ans[-1] = n
    tmp_res = calc_ans(ans)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans_res = tmp_res
print(ans_res)

Buggy submission (verdict: Runtime Error):
import copy
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
ans_res = calc_ans(ans)
for n in etc:
    tmp = copy.copy(ans)
    tmp[-1] = n
    tmp_res = calc_ans(tmp)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans = tmp
        ans_res = tmp_res
print(ans_res)

Buggy submission (verdict: Runtime Error):
import copy
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
ans_res = calc_ans(ans)
for n in etc:
    tmp = copy.copy(ans)
    tmp[-1] = n
    tmp_res = calc_ans(tmp)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans = tmp
        ans_res = tmp_res
print(tmp_res)

Buggy submission (verdict: Runtime Error):
import copy
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    return tmp + (x * x)

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
if N <= K:
    ans = 0
    for i in dic.keys():
        ans += sum(dic[i])
    ans += len(dic.keys()) * len(dic.keys())
    print(ans)
    exit()
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
for n in etc:
    ans_res = calc_ans(ans)
    tmp = copy.copy(ans)
    tmp[-1] = n
    tmp_res = calc_ans(tmp)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans = copy.copy(tmp)
print(tmp_res)

Buggy submission (verdict: Runtime Error):
import copy
N, K = map(int, input().split())
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    tmp += x * x
    return tmp

for i in range(N):
    l = list(map(int, input().split()))
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n <= uniq[K-1][0]:
          break
      etc.append((n, 0))
etc.sort(reverse=True)
for n in etc:
    ans_res = calc_ans(ans)
    tmp = copy.copy(ans)
    tmp[-1] = n
    tmp_res = calc_ans(tmp)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans = copy.copy(tmp)
print(tmp_res)

Buggy submission (verdict: Runtime Error):
import copy
N, K = map(int, input().split())
L = [list(map(int, input().split())) for i in range(N)]
dic = {}
ans = []
uniq = []
etc = []
def calc_ans(ans):
    x = 0
    tmp = 0
    for i in range(K):
      tmp += ans[i][0]
      x += ans[i][1]
    tmp += x * x
    return tmp

for l in L:
    if dic.get(l[0]):
        dic[l[0]].append(l[1])
    else:
        dic[l[0]] = [l[1]]
for t in dic.keys():
    dic[t].sort(reverse=True)
    uniq.append((dic[t][0], t))
uniq.sort(reverse=True)
for i in range(K):
  ans.append((uniq[i][0], 1))
  for n in dic[uniq[i][1]][1:]:
      if n > uniq[K-1][0]:
          etc.append((n, 0))
etc.sort(reverse=True)
for n in etc:
    ans_res = calc_ans(ans)
    tmp = copy.copy(ans)
    tmp[-1] = n
    tmp_res = calc_ans(tmp)
    if ans_res > tmp_res:
        print(ans_res)
        exit()
    else:
        ans = copy.copy(tmp)
print(tmp_res)