Created
August 21, 2018 15:03
-
-
Save pchtsp/15df9cdf579be5634e458c0d1f74d753 to your computer and use it in GitHub Desktop.
The function below takes two nodes (node1, node2) that share the same tree root (which default to the first node in the tree and the last one respectively). It then searches for nodes that are 'in between'. It also determines the order of the two nodes in the tree (which one is first and which one is second).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def get_nodes_between_nodes_in_tree(node1=None, node2=None, only_order=False): | |
""" | |
This procedure searches a tree for all nodes between node1 and node2. | |
In case node1 is None: it should start in the first node | |
If node2 is None: it should end in the last node | |
:param node1: | |
:param node2: | |
:param only_order: we only care about which nodes comes before. | |
Not the nodes in between. | |
:return: (n1, n2, list): the order of the nodes and a list of nodes in between. | |
""" | |
nodes = [] | |
try_to_add_node1 = False | |
try_to_add_node2 = False | |
if node1 is None and node2 is None: | |
raise ValueError("node1 and node2 cannot be None at the same time") | |
if node1 is None: | |
node1 = nd.get_descendant(node2.get_tree_root(), which='first') | |
try_to_add_node1 = True | |
elif node2 is None: | |
node2 = nd.get_descendant(node1.get_tree_root(), which='last') | |
try_to_add_node2 = True | |
else: | |
assert node1.get_tree_root() == node2.get_tree_root(), \ | |
"node {} does not share root with node {}".format(node1.name, node2.name) | |
ancestor = node1.get_common_ancestor(node2) | |
# if one of the nodes is the ancestor: there's no nodes in between | |
if ancestor in [node1, node2]: | |
return node1, node2, [] | |
if try_to_add_node1: | |
nodes.append(node1) | |
if try_to_add_node2: | |
nodes.append(node2) | |
n1ancestors = node1.get_ancestors() | |
n2ancestors = node2.get_ancestors() | |
all_ancestors = set(n1ancestors) | set(n2ancestors) | |
first_node = None | |
second_node = None | |
def is_not_ancestor(node): | |
return node not in all_ancestors | |
for node in ancestor.iter_descendants(strategy='postorder', is_leaf_fn=is_not_ancestor): | |
if first_node is None: | |
if node not in [node1, node2]: | |
continue | |
first_node, second_node = node1, node2 | |
if node == node2: | |
first_node, second_node = node2, node1 | |
if only_order: | |
break | |
continue | |
if node == second_node: | |
break | |
if is_not_ancestor(node): | |
nodes.append(node) | |
return first_node, second_node, nodes |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment