Problem p02850 (python)

Fixed submission:
"""Coloring Edges on Tree
https://atcoder.jp/contests/abc146/tasks/abc146_d
"""
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 5 + 1)

class Node:
    def __init__(self, num):
        self.num = num
        self.edge = set()
        self.edge_color = set()


def dfs(node, c = -1 , p = -1):
    k = 1    
    for e in node.edge:
        if e == p:
            continue
        if k == c:
            k += 1
        # 子のノードを取得
        child = tree[e]

        # 答えを記録
        if a < b:
            ans['{0} {1}'.format(node.num, child.num)] = k
        else:
            ans['{0} {1}'.format(child.num, node.num)] = k
        k += 1

        # 掘る
        dfs(child,k - 1, node.num)


N = int(input())
ans = {}
# colors = {1}

# treeの初期化
tree = [Node(i) for i in range(0, N + 1)]
edges = []
for i in range(N - 1):
    a, b = map(int, input().split())
    tree[a].edge.add(b)
    tree[b].edge.add(a)
    if a < b:
        edges.append('{0} {1}'.format(a, b))
    else:
        edges.append('{0} {1}'.format(b, a))

dfs(tree[1], -1)

# print("edges:{}".format(edges))
# print(max(c))
# print(ans)
count = 0
for node in tree:
    count = max(count, len(node.edge))
print(count)
for edge in edges:
    # print(edge)
    print(str(ans.get(edge)))

Buggy submission (verdict: Runtime Error):
"""Coloring Edges on Tree
https://atcoder.jp/contests/abc146/tasks/abc146_d
"""
import sys
input = sys.stdin.readline


class Node:
    def __init__(self, num):
        self.num = num
        self.edge = []
        self.edge_color = []


def dfs(node, p):
    # print(node.num, p)
    for e in node.edge:
        used_color = node.edge_color
        if e != p:
            # 子のノードを取得
            child = tree[e]
            child_used_color = child.edge_color

            # このnodeと子のノードでまだ使っていない色を特定
            unused_colors = list(set(colors) - set(used_color + child_used_color))
            if len(unused_colors) == 0:
                color_num = max(colors) + 1
                colors.append(color_num)
            else:
                color_num = unused_colors[0]

            # このノードと、子のノードに使用ずみの色をマーク
            node.edge_color.append(color_num)
            child.edge_color.append(color_num)

            # 答えを記録
            ans['{a} {b}'.format(a=node.num, b=child.num)] = color_num
            ans['{a} {b}'.format(a=child.num, b=node.num)] = color_num

            # 掘る
            dfs(child, node.num)


N = int(input())
ans = {}
colors = [1]

# treeの初期化
tree = []
edges = []
tree.append(Node(0))
for i in range(1, N + 1):
    tree.append(Node(i))
for i in range(N - 1):
    inp = input().split()
    edges.append(inp)
    a, b = map(int, inp)
    tree[a].edge.append(b)
    tree[b].edge.append(a)

dfs(tree[1], -1)


print(len(colors))
for e in edges:
    # print(e)
    print(str(ans.get(' '.join(e))))

Buggy submission (verdict: Runtime Error):
"""Coloring Edges on Tree
https://atcoder.jp/contests/abc146/tasks/abc146_d
"""
import sys
input = sys.stdin.readline


class Node:
    def __init__(self, num):
        self.num = num
        self.edge = []
        self.edge_color = []


def dfs(node, p):
    # print(node.num, p)
    for e in node.edge:
        used_color = node.edge_color
        if e != p:
            # 子のノードを取得
            child = tree[e]
            child_used_color = child.edge_color

            # このnodeと子のノードでまだ使っていない色を特定
            unused_colors = list(set(colors) - set(used_color + child_used_color))
            if len(unused_colors) == 0:
                color_num = max(colors) + 1
                colors.append(color_num)
            else:
                color_num = unused_colors.pop()

            # このノードと、子のノードに使用ずみの色をマーク
            node.edge_color.append(color_num)
            child.edge_color.append(color_num)

            # 答えを記録
            ans['{a} {b}'.format(a=node.num, b=child.num)] = color_num

            # 掘る
            dfs(child, node.num)


N = int(input())
ans = {}
colors = [1]

# treeの初期化
tree = []
edges = []
tree.append(Node(0))
for i in range(1, N + 1):
    tree.append(Node(i))
for i in range(N - 1):
    inp = input().split()
    edges.append(inp)
    a, b = map(int, inp)
    tree[a].edge.append(b)
    tree[b].edge.append(a)


dfs(tree[1], -1)


print(len(colors))
for e in edges:
    # print(e)
    print(str(ans.get(' '.join(e))))