Problem p03044 (Python)

Buggy submission (status: Time Limit Exceeded):
N = int(input())
node_array = [list(map(int, input().split())) for i in range(N - 1)]
# print(node_array)
tree_num_array = [-1] * N
# print(tree_num_array)

tree_num_max = -1
tree = {}
last_tree_num = -1
for node in node_array:
    # print(tree)
    i, j, w = node
    i = i - 1
    j = j - 1
    if tree_num_array[i] == -1 and tree_num_array[j] == -1:
        # print(i, j, "added")
        tree_num_max += 1
        tree[tree_num_max] = {i: 0, j: w}
        tree_num_array[i] = tree_num_max
        tree_num_array[j] = tree_num_max
        last_tree_num = tree_num_max
        continue
    elif tree_num_array[i] == -1:
        # print(i, "added")
        tree_num_added = tree_num_array[j]
        tree[tree_num_added][i] = tree[tree_num_added][j] + w
        tree_num_array[i] = tree_num_added
        last_tree_num = tree_num_added
    elif tree_num_array[j] == -1:
        # print(j, "added")
        tree_num_added = tree_num_array[i]
        tree[tree_num_added][j] = tree[tree_num_added][i] + w
        tree_num_array[j] = tree_num_added
        last_tree_num = tree_num_added
    else:
        tree_num_added = tree_num_array[i]
        tree_num_united = tree_num_array[j]
        # print(tree_num_added, tree_num_united, "united")
        tree_value_i = tree[tree_num_added][i]
        tree_value_j = tree[tree_num_united][j]
        for key, value in tree[tree_num_united].items():
            tree[tree_num_added][key] = tree_value_i + w + abs(value - tree_value_j)
            tree_num_array[int(key)] = tree_num_added
        last_tree_num = tree_num_added
# print(tree)
ans_dict = sorted(tree[last_tree_num].items(), key=lambda x: int(x[0]))
# print(ans_dict)
for distance in ans_dict:
    print(distance[1] % 2)


Fixed submission:
N = int(input())
node_array = [list(map(int, input().split())) for i in range(N - 1)]
# print(node_array)
tree_num_array = [-1] * N
# print(tree_num_array)

tree_num_max = -1
tree = {}
last_tree_num = -1
for node in node_array:
    # print(tree)
    i, j, w = node
    i = i - 1
    j = j - 1
    if tree_num_array[i] == -1 and tree_num_array[j] == -1:
        # print(i, j, "added")
        tree_num_max += 1
        tree[tree_num_max] = {i: 0, j: w}
        tree_num_array[i] = tree_num_max
        tree_num_array[j] = tree_num_max
        last_tree_num = tree_num_max
        continue
    elif tree_num_array[i] == -1:
        # print(i, "added")
        tree_num_added = tree_num_array[j]
        tree[tree_num_added][i] = tree[tree_num_added][j] + w
        tree_num_array[i] = tree_num_added
        last_tree_num = tree_num_added
    elif tree_num_array[j] == -1:
        # print(j, "added")
        tree_num_added = tree_num_array[i]
        tree[tree_num_added][j] = tree[tree_num_added][i] + w
        tree_num_array[j] = tree_num_added
        last_tree_num = tree_num_added
    else:
        tree_num_i = tree_num_array[i]
        tree_num_j = tree_num_array[j]
        # print(tree_num_i, tree_num_j, "unite")
        if len(tree[tree_num_i]) > len(tree[tree_num_j]):
            tree_main_num = tree_num_i
            tree_main_value = tree[tree_main_num][i]
            tree_sub_num = tree_num_j
            tree_sub_value = tree[tree_sub_num][j]
        else:
            tree_main_num = tree_num_j
            tree_main_value = tree[tree_main_num][j]
            tree_sub_num = tree_num_i
            tree_sub_value = tree[tree_sub_num][i]

        for key, value in tree[tree_sub_num].items():
            tree[tree_main_num][key] = tree_main_value + w + abs(value - tree_sub_value)
            tree_num_array[int(key)] = tree_main_num
        last_tree_num = tree_main_num
# print(tree)
ans_dict = sorted(tree[last_tree_num].items(), key=lambda x: int(x[0]))
# print(ans_dict)
for distance in ans_dict:
    print(distance[1] % 2)