Traversing through a tree but avoiding certain subtrees

[hyanwong]

A basic question: If I'm traversing through a tskit tree, how can I omit a subtree but carry on looking at the rest of the tree?

The context is: I have some tracked samples, and want to find all the nodes under a focal node in which the number of tracked samples for that node equals the number of samples. I can completely avoid traversing into subtrees it they don't contain any of the tracked samples.

I'm not clear to me how to avoid a subtree during traversal? I'm assuming this is quite a common thing to want to do.

[hyanwong]

Aha - I forgot that I implemented this at

tskit/python/tskit/drawing.py

Line 596 in da094c4

def _postorder_tracked_node_traversal(tree, root, collapse_tracked, key_dict=None):

!

Although the preorder version is much easier:

def preorder_tracked_node_traversal(tree, root):
    # Preorder traversal that only descends into subtrees if they contain
    # a tracked node. If we deliberately specify the virtual root, it should also be returned
    is_virtual_root = root == tree.virtual_root
    if root == tskit.NULL:
        root = tree.virtual_root
    stack = [root]
    while stack:
        u = stack.pop()
        if u != tree.virtual_root or is_virtual_root:
            yield u
        
        if tree.num_children(u) > 0 and tree.num_tracked_samples(u) > 0:
            # Add children in reverse order so they're processed left-to-right
            stack.extend(tree.children(u))

I think it could be helpful to have this code somewhere in the docs, and also to know if there are suitable numba-ified versions available.