Problem p02850 (python)

Fixed submission:
#Edgeの組を拾う時に，辞書が良いかも？
import sys
input = sys.stdin.readline

N = int(input())
E=[[] for _ in range(N)]

index ={}

for i in range(N-1):
    a,b = map(int,input().split())
    E[a-1].append(b-1)
    E[b-1].append(a-1)
    index[(a-1,b-1)]=i
    index[(b-1,a-1)]=i

K=0
#ここに答えを入れていく，
ans=[None]*(N-1)

#dfs,root=-1
#Q=[(node,pare)]
Q=[(0,-1)]
while(Q):
    node,p=Q.pop()
    col=1
    #根の時
    if p==-1:
        for to in E[node]:
            #ansの適切な位置に色情報を入れる
            ans[index[(node,to)]]=col
            col +=1
            K=max(col,K)
            Q.append((to,node))
    #根ではない時
    else:
        for to in E[node]:
            if to==p:
                continue
            #親の色と同じなら1を足す．
            col+=(col==ans[index[(node,p)]]) 
            ans[index[(node,to)]]=col
            col+=1
            K=max(col,K)
            Q.append((to,node))

print(K-1)
print(*ans,sep="\n")

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

N = int(input())
A = [None]*N
B = [None]*N
K = 0

adj_list=[[] for _ in range(N)]
c_mat = [[-1]*N for _ in range(N)]
c_min_list=[-1]*N #A[i]に隣接する中で使っている色の最小値をメモ

for i in range(N-1):
    A[i],B[i] = map(int,input().split())
    adj_list[A[i]-1].append(B[i]-1)
    adj_list[B[i]-1].append(A[i]-1)
    
#着色
def dfs(node_num,p_num,p_col):
    for i in range(len(adj_list[node_num])):
        to=adj_list[node_num][i]
        #親以外につながる辺を塗る
        if to!=p_num:
            #色を決める
            c_to = c_min_list[node_num]+1
            if c_to==p_col:
                c_to+=1
            c_min_list[node_num] = c_to
            c_mat[node_num][to] = c_to
            c_mat[to][node_num] = c_to
            dfs(to,node_num,c_to)



#計算，出力
dfs(0,-1,N)
K = max(c_min_list)+1
print(K)
for i in range(N-1):
    print(c_mat[A[i]-1][B[i]-1]+1)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

N = int(input())
A = [None]*N
B = [None]*N
K = 0

adj_list=[[] for _ in range(N)]
c_list = [[] for _ in range(N)]
c_min_list=[-1]*N #A[i]に隣接する中で使っている色の最小値をメモ

for i in range(N-1):
    A[i],B[i] = map(int,input().split())
    adj_list[A[i]-1].append(B[i]-1)
    adj_list[B[i]-1].append(A[i]-1)
    
#幅優先で着色
def bfs(node_num,p_num,p_col):
    for i in range(len(adj_list[node_num])):
        to=adj_list[node_num][i]
        #親以外につながる辺を塗る
        if to!=p_num:
            #色を決める
            c_to = c_min_list[node_num]+1
            if c_to==p_col:
                c_to+=1
            c_min_list[node_num] = c_to
            c_list[node_num].append((to,c_to))
            c_list[to].append((node_num,c_to))
            bfs(to,node_num,c_to)



#計算，出力
bfs(0,-1,N)
K = max(c_min_list)+1
print(K)
for i in range(N-1):
    for j in range(N-1):
        if c_list[A[i]-1][j][0]==B[i]-1:
            print(c_list[A[i]-1][j][1]+1)
            break

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

N = int(input())
A = [None]*N
B = [None]*N
K = 0

adj_list=[[] for _ in range(N)]
c_list = [[] for _ in range(N)]

for i in range(N-1):
    A[i],B[i] = map(int,input().split())
    adj_list[A[i]-1].append(B[i]-1)
    adj_list[B[i]-1].append(A[i]-1)
    
#幅優先で着色
def bfs(node_num,p_num,p_col):
    global K
    for i in range(len(adj_list[node_num])):
        to=adj_list[node_num][i]
        #親以外につながる辺を塗る
        if to!=p_num:
            #色をiで振っていくが，もし親の色がiよりも小さい場合，色を１つ進める
            c_to=i
            if i>=p_col:
                c_to+=1
            if c_to>K:
                K=c_to
            c_list[node_num].append(c_to)
            bfs(to,node_num,c_to)
        else:
            #indexを揃えるため-1で埋める？
            c_list[node_num].append(-1)

#計算，出力
bfs(0,-1,-1)
print(K)
for i in range(N-1):
    j=adj_list[A[i]-1].index(B[i]-1)
    print(c_list[A[i]-1][j])