Binary Search Tree - node count method with an incorrect output

Question

So basically, as you can see from the function definition this function is supposed to determine the node type of a Binary Search Tree, I am not getting errors, however I don't think my output is correct, I have a auxiliary method to recursively search through the BST and check to see wether it's a zero, one or two. If I input this BST below:

    22
   /  \
  12   30
 / \   / \
8  20 25 40

it returns 0,0,1 but I don't think that's correct, shouldn't it be returning 4,0,3? since 22,12 and 30 are two's as they have two child nodes and 8,20,25,and 40 are zero's as they point to leafs. Any help would be greatly appreciated!

Here is my code:

def node_counts(self):
    """
    ---------------------------------------------------------
    Returns the number of the three types of nodes in a BST.
    Use: zero, one, two = bst.node_counts()
    -------------------------------------------------------
    Returns:
        zero - number of nodes with zero children (int)
        one - number of nodes with one child (int)
        two - number of nodes with two children (int)
    ----------------------------------------------------------
    """
    zero, one, two = self.node_counts_aux(self._root)
    return zero, one, two

    return

def node_counts_aux(self, node):
    zero = 0
    one = 0
    two = 0

    if node is None:
        return zero, one, two

    else:
        self.node_counts_aux(node._left)
        print(node._value)
        self.node_counts_aux(node._right)

        if node._left is not None and node._right is not None:
            two += 1
        elif (node._left is not None and node._right is None) or (node._left is None and node._right is not None):
            one += 1
        else:
            zero += 1

    return zero, one, two

I am currently doing an inorder traversal, I am expecting this 4,0,3 instead of this 0,0,1

Answer 1

A common mistake with recursion is forgetting to use returned values. You need to pass them up the call tree for them to "count" at the top level:

def node_counts_aux(self, node):
    zero = 0
    one = 0
    two = 0

    if node is None:
        return zero, one, two

    z, o, t = self.node_counts_aux(node._left)
    zero += z
    one += o
    two += t
    z, o, t = self.node_counts_aux(node._right)
    zero += z
    one += o
    two += t

    if node._left and node._right:
        two += 1
    elif node._left or node._right:
        one += 1
    else:
        zero += 1

    return zero, one, two

That said, generally I'd prefer an internal function than a helper function and use a list instead of distinct variables:

def node_counts(self):
    counts = [0, 0, 0]

    def traverse(node):
        if not node:
            return

        traverse(node._left)
        traverse(node._right)

        if node._left and node._right:
            counts[2] += 1
        elif node._left or node._right:
            counts[1] += 1
        else:
            counts[0] += 1

    traverse(self._root)
    return tuple(counts)

Much cleaner; less data to have to pass around. It's safe to assume truthy values for nodes.

Note also that the traversal order is irrelevant.

Runnable example on your tree:

from collections import namedtuple


def count_node_types(root):
    counts = [0, 0, 0]

    def traverse(node):
        if not node:
            return

        traverse(node.left)
        traverse(node.right)

        if node.left and node.right:
            counts[2] += 1
        elif node.left or node.right:
            counts[1] += 1
        else:
            counts[0] += 1

    traverse(root)
    return tuple(counts)


if __name__ == "__main__":
    Node = namedtuple("Node", "val left right", defaults=[None] * 3)
    r""" 22
       /  \
      12   30
     / \   / \
    8  20 25 40 """
    root = Node(22, Node(12, Node(8), Node(20)), Node(30, Node(25), Node(40)))
    print(count_node_types(root))  # => (4, 0, 3)

... But it's better to pick a test case that ensures nodes with 1 child are counted properly, so I'd remove Node(40) to make sure this returns (3, 1, 2).