Another Simple Tree in Python

I’ve posted before about creating a tree in Python, but I like this implementation better. It uses a nested class to represent the nodes of the tree, and an interesting construction (line 11) that is a result of that nested class. Also, I do a simple pre-order traversal. I’ll flesh this guy out in later posts.

class Tree:
    
    class Node:
        def __init__( self, data ):
            self.parent = None
            self.data = data
            self.children = []
        def addChild( self, data ):
            # I've actually never used this
            # construction in a class before!
            node = Tree.Node( data )
            node.parent = self
            self.children.append( node )
        def getChild( self, data ):
            for child in self.children:
                if child.data == data:
                    return child
            return None
        def countAllChildren( self ):
            """
            pre-order traversal to count all
            of the children below this node
            """
            # for this node, get a count
            # of the immediate children
            count = len( self.children )
            for child in self.children:
                # then call this method for each child
                # and increment the total count in the
                # scope of this node
                count += child.countAllChildren()
            # then return that count
            return count
        def __repr__( self ):
            return self.data
        def __len__( self ):
            return self.countAllChildren()
    
    def __init__( self, data ):
        self.root = self.Node( data )

And some (sort of) tests this time!

t = Tree("/")
t.root.addChild("usr/")
t.root.addChild("bin/")
t.root.getChild("usr/").addChild("local/")

assert t.root.parent == None
assert len( t.root.children ) == 2
assert t.root.getChild("usr/").parent == t.root
assert len( t.root.getChild("usr/").children ) == 1
assert t.root.countAllChildren() == 3
assert len( t.root ) == 3

Leave a Reply

Your email address will not be published. Required fields are marked *