Skip to content

Instantly share code, notes, and snippets.

@pchtsp
Created August 21, 2018 15:03
Show Gist options
  • Save pchtsp/15df9cdf579be5634e458c0d1f74d753 to your computer and use it in GitHub Desktop.
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).
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