binary_tree_node_count.py

# Given a complete binary tree (every level full, except possibly the last level,
# but full left to right in any case; no gaps):
# count the number of nodes in the tree

# 1. go down left side of tree -> lg N time to do; number of nodes will be between
#       2 ** (leftmost depth - 1) to 2 ** (leftmost depth) - 1
#    where depth includes the root:
# Size = 2 ** 2 (4)   -> 2 ** 3 - 1 (7)
# Depth = 3
#
#             X             
#       X           X          
#   X     X      X     X  
#  X X   X _    _ _   _ _
# 
# 2. root.left.right.right.right... 
#    now recurse with root.left / root.right as root
#
# count all nodes in O(lg N x lg N)

class Node(object):
    pass


def number_of_nodes(tree):
    def dfs(node, d, mode):
        if not node:
            return d

        if mode == 'left':
            return dfs(node.left, d + 1, 'left')
        else:
            return dfs(node.right, d + 1, 'right')


    def subtree_bin_search_count(node, max_depth, node_depth, count):
        if node_depth == max_depth:
            if node:
                return count + 1
            else:
                return count
        if dfs(node.left, node_depth, 'right') == max_depth:
            # everything in node.left is full
            count += 2 ** (max_depth - node_depth - 1)
            return subtree_bin_search_count(node.right, max_depth, node_depth + 1, count)
        else:
            return subtree_bin_search_count(node.left, max_depth, node_depth + 1, count)
        
    if not tree:
        return 0

    hleft = dfs(tree, 0, 'left')
    hright = dfs(tree, 0, 'right')  # == hleft or == hleft - 1

    if hleft == hright:
        return (2 ** hleft) - 1
    return subtree_bin_search_count(tree, hleft, 1, (2 ** hright) - 1)

#             X             
#       X           X          
#   X     X      X     X  
#  X X   X _    _ _   _ _
# 
# hleft = 4
# hright = 3
# subtree_bin_search_count(tree, 4, 1, 7)
#   4 != 1
#   dfs(tree.left, 1, 'right') == 3 != 4
#   return subtree_bin_search_count(tree.left, 4, 2, 7)
#   4 != 2
#   dfs(tree.left.left, 2, 'right') == 4 == 4
#   count += 2
#   return subtree_bin_search_count(tree.left.right, 4, 3, 7 + 2 == 9)
#       4 != 3
#       dfs(tree.left.right.left, 3, 'right') == 3 != 4
#       return subtree_bin_search_count(tree.left, 4, 4, 9)
#           4 == 4: return 9 + 1 == 10