Counting the nodes in a Binary Search Tree in Python

Question

I'm fairly new to programming and I want to screw around with some Binary Search Trees. I want to make a function that counts the number of nodes in the tree recursively, however, when I run my function it doesn't seem to work and it keeps returning 'none' as if there is nothing in my tree. Could anyone help me find the problem here?

This is my TreeNode class:

class TreeNode(object):

    def __init__(self, data = None, left=None, right=None):
        self.item = data
        self.left = left
        self.right = right

    def __str__(self):
        return str(self.item)

This is my main function, I trimmed most of it down just so we can get to the problem referring to the counting of the nodes.

from TreeNode import TreeNode


class BST(object):

    #------------------------------------------------------------

    def __init__(self):

        """create empty binary search tree
        post: empty tree created"""

        self.root = None

def treeSize(self, root, size = 0):

        if root is None:
            return -1

        if root is not None:
            size += 1
            if root.left is not None:
                self.treeSize(root.left, size)
            if root.right is not None:
                self.treeSize(root.right, size)

This is the code I use to test out my function:

from BinarySearchTree import BST
from TreeNode import TreeNode

tree = TreeNode(4, TreeNode(2, TreeNode(1), TreeNode(3)), TreeNode (7, TreeNode(6),TreeNode(8)))

a = BST()

print(a.postOrder(tree))
print(a.treeSize(tree))

When I called the 'print(a.treeSize(tree))' It just returns 'none' and not '7' like it should.

Answer 1

You can also do it the good old recursive way:

def treeSize(self, root):

    if root is None:
        return 0

    if root is not None:
        return 1 + self.treeSize(root.left) + self.treeSize(root.right)

Jonathan's answer is nice as well.

Answer 2

I see. You think size is going to get updated in the called functions. It won't as it's local to each function. You could call global on it, but that's not optimal.

You could set it as a member variable (don't actually do it this way):

def __init__(self):
   ...
   self.size = 0

def treeSize(self,...):
   ...
   self.size += 1
   ...
   return self.size

but the obvious bug is self.size will double every time treeSize is called. You can fix that too, but let's use patterns we know and love. Do it the good old recursive way like VHarisop wrote.