Problem p03148 (python)

Fixed submission:
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k != 0:
        if a[A[i][0]] ==0:
            a[A[i][0]] = 1
            eat.append(A[i][1])
            k -= 1
        else:
            noteat.append(A[i][1])
    else:
        noteat.append(A[i][1])
        
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat[k:])+1)):
    ans[i] = ans[i-1]-(pow(K-k-(i-1),2)+eat[-i]) +(pow(K-k-i,2)+noteat[k+i-1])

print(max(ans))

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat))+2)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat)+1)):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat))+1)]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = [0 for i in range(min(K-k,len(noteat)))]
ans[0] = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
#    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
    ans[i] = ans[i-1]-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
#    else:
#        break
print(max(ans))

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
eat,noteat =[],[]
a = [0 for i in range(N+1)]
k = K
for i in range(N):
    if k == 0:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans = sum(eat)+ sum(noteat[:k])+ pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
        ans = ans-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
    else:
        break
print(ans)

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Mon Jan 21 20:39:09 2019

@author: Yamazaki Kenichi
"""

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

A.sort(reverse = True,key=lambda x:x[1])
ans = 0
eat,noteat =[],[]
a = [0 for i in range(N)]
k = K
for i in range(N):
    if k == -1:
        break    
    if a[A[i][0]] ==0:
        a[A[i][0]] = 1
        eat.append(A[i][1])
        k -= 1
    else:
        noteat.append(A[i][1])
ans += sum(eat)
ans += sum(noteat[:k])
ans += pow(K-k,2)

for i in range(1,min(K-k,len(noteat))):
    if pow(K-k,2)+eat[-i] < pow(K-k-i,2)+noteat[k:][i-1]:
        ans = ans-(pow(K-k,2)+eat[-i]) +(pow(K-k-i,2)+noteat[k:][i-1])
    else:
        break
print(ans)

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 21:13:02 2019

@author: Yamazaki Kenichi
"""

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

A.sort(key=lambda x: x[1])
kinds = 0
t = 0
for i in range(N):
    if t != A[i][0]:
        kinds += 1
        t = A[i][0]
c = K - kinds if K - kinds >=0 else 0
d = K-c
temp_t = pow(K-c,2)
i = N-1
counta =0
temp_d=0
a = [0 for i in range(N+1)]
b,b2 =[],[]
while counta < K:
    if c-a[A[i][0]] >= 0:
        temp_d += A[i][1]
        c = c-1 if a[A[i][0]] >= 1 else c
        a[A[i][0]] += 1
        counta+=1
        b2.append(A[i][1])
    else:
        b.append(A[i][1])
    i -=1
ans = [0 for i in range(len(b)+1)]
ans[0] = temp_t+temp_d
#print(ans)
for i in range(len(b)):
    ans[i+1] = ans[i] +b[i]-b2[-1-i]-(pow(d,2)-pow(d-i-1,2))
print(max(ans))

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
"""
Created on Sun Jan 20 21:13:02 2019

@author: Yamazaki Kenichi
"""

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

A.sort(key=lambda x: x[1])
kinds = 0
t = 0
for i in range(N):
    if t != A[i][0]:
        kinds += 1
        t = A[i][0]
c = K - kinds if K - kinds >=0 else 0
d = K-c
temp_t = pow(K-c,2)
i = N-1
counta =0
temp_d=0
a = [0 for i in range(N+1)]
b,b2 =[],[]
while counta < K:
    if c-a[A[i][0]] >= 0:
        temp_d += A[i][1]
        c = c-1 if a[A[i][0]] >= 1 else c
        a[A[i][0]] += 1
        counta+=1
        b2.append(A[i][1])
    else:
        b.append(A[i][1])
    i -=1
ans = temp_t+temp_d
#print(ans)
for i in range(len(b)):
    if b[i]-b2[-1-i]>pow(d,2)-pow(d-i-1,2):
        ans = ans +b[i]-b2[-1-i]-(pow(d,2)-pow(d-i-1,2))
    else:
        break
print(ans)