Problem p02850 (python)

Fixed submission:
class ColorMgmt:
    from heapq import heappop, heappush
    def __init__(self, n):
        self.colors_hash = [set() for _ in range(n)]
        self.next_colors = [1] * n
    
    def set_color(self, v, c):
        self.colors_hash[v].add(c)
        while self.next_colors[v] in self.colors_hash[v]:
            self.next_colors[v] += 1
    
    def get_color(self, v):
        return self.next_colors[v]


N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]
edge = [[] for _ in range(N)]
edge_cnt = [0] * N
for idx, (a, b) in enumerate(ab):
    a, b = a-1, b-1
    edge[a].append([b, a, 0, idx])
    edge_cnt[a] += 1
    edge_cnt[b] += 1
K = max(edge_cnt)

next_list = []
now_list = [0]
color_mgmt = ColorMgmt(N)
while len(now_list) > 0:
    for v in now_list:
        for idx, nextv in enumerate(edge[v]):
            if v == nextv[0]:
                continue
            colors1 = color_mgmt.get_color(nextv[1])
            colors2 = color_mgmt.get_color(nextv[0])
            color = max(colors1, colors2)
            color_mgmt.set_color(nextv[1], color)
            color_mgmt.set_color(nextv[0], color)
            edge[v][idx][2] = color
            next_list.append(nextv[0])
    
    now_list = list(next_list)
    next_list = []

ret = []
for edges in edge:
    for e in edges:
        ret.append(e)
ret.sort(key=lambda x:x[3])
print(K)
for r in ret:
    print(r[2])

Buggy submission (verdict: Runtime Error):
N = int(input())
edge = [[] for _ in range(N)]
for i in range(N-1):
    a, b = map(int, input().split())
    a, b = a-1, b-1
    edge[a].append(b)


color = [0] * N

def dfs(v, c):
    color[v] = c
    now = 0
    for nv in edge[v]:
        now += 1
        if now == c:
            now += 1
        dfs(nv, now)

dfs(0, 0)
color[0] = max(color)
print(*color, sep="\n")

Buggy submission (verdict: Runtime Error):
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]
edge = [[] for _ in range(N)]
for idx, (a, b) in enumerate(ab):
    a, b = a-1, b-1
    edge[a].append([b, a, 0, idx])

next_list = []
now_list = [0]
colors = [0] * (N-1)
color_list = set()
now_color = 2
max_color = 0
while len(now_list) > 0:
    for v in now_list:
        for idx, nextv in enumerate(edge[v]):
            if v == nextv[0]:
                continue
            colors[nextv[0]] += 1
            edge[v][idx][2] = colors[nextv[0]]
            max_color = max(max_color, now_color-1)
            next_list.append(nextv[0])
    
    now_list = list(next_list)
    next_list = []

ret = []
for ed in edge:
    for e in ed:
        ret.append(e)
ret.sort(key=lambda  x:x[3])
print(max_color)
for ed in ret:
    print(ed[2])

Buggy submission (verdict: Runtime Error):
N = int(input())
ab = [list(map(int, input().split())) for _ in range(N-1)]
edge = [[] for _ in range(N)]
for a, b in ab:
    a, b = a-1, b-1
    edge[a].append(b)

next_list = []
now_list = [0]
colors = [0] * N
colors[0] = 1
color_list = set()
now_color = 2
while len(now_list) > 0:
    if not 1 in color_list:
        now_color = 1
    color_list = set()
    for v in now_list:
        for nextv in edge[v]:
            if v == nextv:
                continue
            colors[nextv] = now_color
            color_list.add(now_color)
            now_color += 1
            next_list.append(nextv)
    
    now_list = list(next_list)
    next_list = []
            
print('\n'.join(colors))