Problem p03148 (python)

Fixed submission:
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N+1)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    dt[i][1]-=1
    ans+=dt[i][0]
    if gr[dt[i][1]]==[]:
        x+=1
    gr[dt[i][1]].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    
inf=10**19
heapq.heappush(single,[inf,N])#絶対に選ばないものを入れておく
heapq.heappush(multi,[inf,N])
    
    
    
for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:#種類ボーナスも含めて小さい方を取り除く
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        tt=b[1]
        num_i[tt]+=1
        if num_i[tt]==len(gr[tt])-1:
            heapq.heappush(single,gr[tt][-1])
        else:
            heapq.heappush(multi,gr[tt][num_i[tt]])
                                       
print(ans)

Buggy submission (verdict: Runtime Error):
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N+1)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    dt[i][1]-=1
    ans+=dt[i][0]
    if gr[dt[i][1]]==[]:
        x+=1
    gr[dt[i][1]].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    
inf=10**19
heapq.heappush(single,[inf,N])#絶対に選ばないものを入れておく
heapq.heappush(multi,[inf,N])
    
    
    
for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:#種類ボーナスも含めて小さい方を取り除く
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        tt=b[1]
        num_i[tt]+=1
        if num_i[tt]==len(gr[tt])-1:
            heapq.heappush(single,gr[tt][-1])
        else:
            heapq.heappush(multi,gr[tt][num_i[tt]])
                                       
print(ans)a

Buggy submission (verdict: Runtime Error):
#残り1種類のグループとそうでないグループに分けておく，残り1種類になったら移動
#判定のため，各tごとにdのリスト（ソート済み）と，dのリストを参照する番号を持っておく

import heapq

N,K=map(int,input().split())
dt=[[0,0] for _ in range(N)]
gr=[[] for _ in range(N)]#tの種類ごとにまとめる
x=0#種類数
ans=0#先に全部食べたことにして，いらないものを除く

for i in range(N):
    dt[i][1],dt[i][0]=map(int,input().split())
    ans+=dt[i][0]
    if gr[dt[i][1]-1]==[]:
        x+=1
    gr[dt[i][1]-1].append(dt[i])
ans+=x*x

single=[]
multi=[]
heapq.heapify(single)
heapq.heapify(multi)

num_i=[0]*N#各tに対してどこを見るか

for i in range(N):
    gr[i].sort()
    if gr[i]!=[]:
        if len(gr[i])==1:
            heapq.heappush(single,gr[i][0])
        else:
            heapq.heappush(multi,gr[i][0])
    

for i in range(N-K):
    a=heapq.heappop(single)
    b=heapq.heappop(multi)
    if a[0]+(2*x-1)<=b[0]:
        ans-=a[0]+(2*x-1)
        x-=1
        heapq.heappush(multi,b)
    else:
        ans-=b[0]
        heapq.heappush(single,a)
        num_i[b[1]-1]+=1
        if num_i[b[1]-1]==len(gr[b[1]-1])-1:
            heapq.heappush(single,gr[b[1]-1][-1])
        else:
            heapq.heappush(multi,gr[b[1]-1][num_i[b[1]-1]])
                                       
print(ans)