Prune tree branches where ending nodes meet some criteria

Question

I have a TreeNode class as follows:

public class TreeNode
{
    public enum NodeType { Root,Element, Category}
    public TreeNode()
    {
        Children = new List<TreeNode>();
    }
    public List<TreeNode> Children { get; set; }
    public string Name { get; set; }
    public NodeType Type { get; set; }
}

And I have a BuildTree method that populates the tree structure setting the node Name and type for each node (e.g. The first node will be set to Root Type, and child nodes could be either Element or Category types).

I'm looking for an efficient way (perhaps using LINQ) to prune tree branches where the ending node type is not of type Category. Any ideas on how to achieve that?

Here a visual example:

Before:

Root
|
|_ NodeA (Element)
|_ Node B (Element)
|  |_ Node B.1 (Category)
|_ Node C (Element)
|  |_ Node C.1 (Element)
|    |_Node C.1.1 (Category)
|_ Node D (Element)
   |_Node D.1 (Element)

After:

Root
|
|_ Node B (Element)
|  |_ Node B.1 (Category)
|_ Node C (Element)
   |_ Node C.1 (Element)
     |_Node C.1.1 (Category)

Thanks in advance!

Answer 1

First we need to describe how to do a depth first aggregation of a tree:

//seed: leafNode -> accumulate
//func: interiorNode, children accumulates -> accumulate
public static TAccumulate Aggregate<TAccumulate>(
    this TreeNode root,
    Func<TreeNode, TAccumulate> seed,
    Func<TreeNode, List<TAccumulate>, TAccumulate> func)
{
    if (root.Children == null || !root.Children.Any())
        return seed(root);
    return func(root, children.Select(child => child.Aggregate(seed, func)).ToList());
}

Then we can use that to make a modification such as the one you requested:

tree.Aggregate(
    leaf => new
    {
        Keep = leaf.NodeType == TreeNode.NodeType.Category,
        Node = leaf
    },
    (node, children) =>
    {
        var removal = new HashSet<TreeNode>(
            children.Where(child => !child.Keep).Select(child => child.Node));
        node.Children.RemoveAll(removal.Contains);
        return new
        {
           Keep = children.Any(child => child.Keep),
           Node = node
        };
    });

This gives you a concise, declarative way to describe what transform you want applied to your tree without getting into the details of the traversal in the description of the transform.

Answer 2

It's not that complex; you just need a recursive traversal of the tree. The base case is that the node itself is a category. Then just call the function on each of the children and keep those that have a category among their children.

/// <summary>
/// Removes all child nodes that do not have at least one Category as a child.
/// </summary>
/// <param name="node">The node to alter the children of.</param>
/// <returns>True if there is at least one Category in this subtree</returns>
public static bool RemoveEmptyElements(TreeNode node)
{
    if (node.Type == TreeNode.NodeType.Category)
        return true;
    node.Children = node.Children.Where(child => RemoveEmptyElements(child))
            .ToList();
    return node.Children.Any();
}