from collections import defaultdict
def make_rooted_tree(list_of_lists):
element_dict = defaultdict(list)
root = 0
ret = []
print("Input: {}".format(list_of_lists))
for i in range(0, len(list_of_lists)):
element_dict[list_of_lists[i][0]].append(list_of_lists[i][1])
# find root and candidates to delete
# find root node
parents = list(element_dict.keys())
children = list(element_dict.values())
children = [item for sublist in children for item in sublist]
children = list(set(children))
print("List of parent nodes: {}".format(parents))
print("List of children nodes: {}".format(children))
# case A: if node in parents but not in children, it's a root
for value in parents:
if value not in children:
root = value
break
# case B: if root cannot be found above, choose the node with the most children
if root == 0:
max_children = 0
for key in element_dict:
number_of_chidlren = len(element_dict[key])
if number_of_chidlren > max_children:
max_children = number_of_chidlren
root = key
print("Root node is {}".format(root))
# find candidate edges to delete
children_of_root = element_dict[root]
print(children_of_root)
# if both nodes of an edge are children of root, then it should be removed
for key in element_dict:
if key in children_of_root:
list_of_key_children = element_dict[key]
for child in list_of_key_children:
if child in children_of_root:
ret.append([key, child])
# if root is a child of a node, edge [node,root] should be removed
for key in element_dict:
if root in element_dict[key]:
ret.append([key, root])
return ret
if __name__ == "__main__":
a = [[1, 2], [1, 3], [2, 3]]
b = [[1, 2], [2, 3], [3, 4], [4, 1], [1, 5]]
deleted_edge_a = make_rooted_tree(a)
print("Edge to be deleted: {} \n ==========================================================".format(deleted_edge_a))
deleted_edge_b = make_rooted_tree(b)
print("Edge to be deleted: {} \n ==========================================================".format(deleted_edge_b))