cut-the-tree-iterative.py

node_sums = {}
edge_dict = {}
seen = set()



def get_forward_branches(node):

    node_edges = edge_dict[node]
    retlist = []

    for target_node in node_edges:
        if target_node not in seen:
            retlist.append(target_node)

    return retlist



def has_forward_branches(node):

    for sub_node in edge_dict[node]:
        if sub_node not in seen:
            return True

    return False



def get_sum(node, node_values):

    '''if node not in seen:
        seen.add(node)'''

    if node in node_sums:
        return node_sums[node]

    retval = node_values[node]

    nodes_to_visit = []

    if has_forward_branches(node):
        nodes_to_visit.extend(get_forward_branches(node))

        for target_node in nodes_to_visit:

            '''if target_node not in seen:
                seen.add(target_node)'''

            if target_node in node_sums:
                retval += node_sums[target_node]

            else:
                retval += node_values[target_node]

                if has_forward_branches(target_node):
                    nodes_to_visit.extend(get_forward_branches(target_node))
                else:
                    node_sums[target_node] = node_values[target_node]


    node_sums[node] = retval
    #print(nodes_to_visit)
    #print(node_sums)
    return retval



def cutTheTree(node_values, edges):

    node_values_sum = sum(node_values)
    half_sum = node_values_sum / 2
    req_sum = 0
    global edge_dict
    global node_sums


    for edge in edges:

        if edge[0] - 1 in edge_dict:
            edge_dict[edge[0] - 1].append(edge[1] - 1)

        else:
            edge_dict[edge[0] - 1] = [edge[1] - 1]


        if edge[1] - 1 in edge_dict:
            edge_dict[edge[1] - 1].append(edge[0] - 1)

        else:
            edge_dict[edge[1] - 1] = [edge[0] - 1]


    nodes_to_visit = [0]

    for node in nodes_to_visit:
        seen.add(node)
        if has_forward_branches(node):
            nodes_to_visit.extend(get_forward_branches(node))

    nodes_to_visit.reverse()
    seen.clear()
    #print(nodes_to_visit)

    for node in range(len(node_values)):
        seen.add(node)

    for node in nodes_to_visit:
        seen.remove(node)
        current_node_sum = get_sum(node, node_values)

        if abs(half_sum - current_node_sum) < abs(half_sum - req_sum):
            req_sum = current_node_sum


    return int((2 * abs(half_sum - req_sum)))