Jul 02 2013

Performing a DFS over a Rooted Tree with a C# foreach Loop

Categories: C#

I’ve been working with trees a lot lately and one thing that occurs quite frequently in my problem space is the action of performing a depth-first search (DFS) or a 'pre-order traversal’.  In a DFS, you visit the root first and then search deeper into the tree visiting each node and then the node’s children.
 

Since this is so common, I’d like a routine that helps me do that very simply.  What I’d really like to do is simply do a foreach loop over the tree and with each iteration pull out the next node that would satisfy a DFS.  This is where the C# yield keyword comes into play.

The following snippet presents a TreeNode class which is just a generic node in a tree that has N children:

public class TreeNode
{
    public TreeNode()
    {
        Children = new List<TreeNode>();
    }

    public TreeNode Parent { get; set; }
    public List<TreeNode> Children { get; private set; }
    public bool IsLeaf { get { return Children.Count == 0; } }
    public bool IsRoot { get { return Parent == null; } }

    public IEnumerator<TreeNode> GetEnumerator()
    {
        yield return this;
        foreach (TreeNode child in Children)
        {
            IEnumerator<TreeNode> childEnumerator = child.GetEnumerator();
            while (childEnumerator.MoveNext())
                yield return childEnumerator.Current;
        }
    }
}

 

Most of the functionality is in the GetEnumerator() method.  When a class in C# has a method called GetEnumerator() that returns a type IEnumerator or IEnumerator<T>, you’ll be able to use that type in a foreach loop.  Also notice that the type returned by my method does not return IEnumerator<TreeNode> but rather just TreeNode.  The C# compiler will fill in the proper details automatically.

In order to do a DFS, we first visit the node itself by returning it (this) in a yield statement.  When the iterator executes the next loop, the GetEnumerator() method remembers the execution state of the method and picks up right where it left off.  At that point, we enter into a loop over the Children collection.  We basically repeat the enumerator process over this collection manually by doing a while loop where we exit the loop if no more children exist.  In the body of this while loop, we execute another return within a yield, returning each child node.

To use this enumerator to perform a DFS, you’d foreach over the root node of your tree (I’m assuming a singly rooted-tree here):

foreach (TreeNode node in root)
{
  // perform action on node in DFS order here
}

No Comments