Returns the size of the largest binary subtree that is complete, in Python - variant for exercise 9.1 on EPI (Elements of Programming Interviews)

find_largest_binary_tree_complete.py

# Returns the size of the largest binary subtree that is complete, in Python - variant for exercise 9.1 on EPI (Elements of Programming Interviews)
# The gist of the solution is to keep track of the subtree's current number of nodes, current height and maximum height until that point.
# With the current number of nodes and height, one can calculate the root's own number of nodes and height via its direct childs respective information,
#   taking into to consideration the relation between the child heights and if they are perfect subtrees or not.
# Solution is O(n) time complexity and O(h) space complexity (function call stack corresponds from the root through the unique path to the current node)

from collections import namedtuple


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

# number of nodes for a perfect tree of the given height
def max_nodes_per_height(height: int) -> int:
    return 2**(height + 1) - 1


def height_largest_complete_subtree(root: BTN) -> int:
    CompleteInformation = namedtuple('CompleteInformation', ['height', 'num_nodes', 'max_height'])
    
    def height_largest_complete_subtree_aux(root: BTN) -> CompleteInformation:
        if (root is None):
            return CompleteInformation(-1, 0, 0)
        
        left_complete_info = height_largest_complete_subtree_aux(root.left)
        right_complete_info = height_largest_complete_subtree_aux(root.right)
        
        left_height = left_complete_info.height
        right_height = right_complete_info.height
        
        if (left_height == right_height):
            if (left_complete_info.num_nodes == max_nodes_per_height(left_height)):
                new_height = left_height + 1
                new_num_nodes = left_complete_info.num_nodes + right_complete_info.num_nodes + 1
                return CompleteInformation(new_height,
                                           new_num_nodes,
                                           max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                                          )
            else:
                new_height = left_height
                new_num_nodes = max_nodes_per_height(left_height)
                return CompleteInformation(new_height,
                                           new_num_nodes,
                                           max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                                          )
        elif (left_height > right_height):
            if (max_nodes_per_height(right_height) == right_complete_info.num_nodes):
                new_height = right_height + 2
                new_num_nodes = min(left_complete_info.num_nodes, max_nodes_per_height(right_height + 1)) + right_complete_info.num_nodes + 1
                return CompleteInformation(new_height,
                               new_num_nodes,
                               max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                              )
            else:
                new_height = right_height + 1
                new_num_nodes = max_nodes_per_height(right_height) + right_complete_info.num_nodes + 1
                return CompleteInformation(new_height,
                               new_num_nodes,
                               max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                              )
        
        elif (left_height < right_height):
            if (left_complete_info.num_nodes == max_nodes_per_height(left_height)):
                new_height = left_height + 1
                new_num_nodes = left_complete_info.num_nodes + max_nodes_per_height(left_height) + 1
                return CompleteInformation(new_height,
                           new_num_nodes,
                           max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                          )
            else:
                new_height = left_height
                new_num_nodes = (max_nodes_per_height(left_height - 1) * 2) + 1
                return CompleteInformation(new_height,
                           new_num_nodes,
                           max(new_height, max(left_complete_info.max_height, right_complete_info.max_height))
                          )
            
    return height_largest_complete_subtree_aux(root).max_height
        
    
tree = BTN('A', BTN('B', BTN('D', BTN('F')), BTN('E')), BTN('C'))
print(height_largest_complete_subtree(tree))

tree = BTN('A', None, BTN('B'))
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D', BTN('G', BTN('K'), BTN('L')), BTN('H')), BTN('E', BTN('I'), BTN('J')))
right = BTN('C', BTN('F'))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D', BTN('G', BTN('K'), BTN('L')), BTN('H')), BTN('E', BTN('I'), BTN('J')))
right = BTN('C', BTN('F'), BTN('CA'))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D', BTN('G', BTN('K'), BTN('L')), BTN('H')), BTN('E', BTN('I'), BTN('J')))
right = BTN('C', BTN('F', BTN('CB'), BTN('CC')), BTN('CA', BTN('CD'), BTN('CE')))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D', BTN('G', BTN('K'), BTN('L')), BTN('H')), BTN('E', BTN('I'), BTN('J')))
right = BTN('C', BTN('F', BTN('CB'), BTN('CC')), BTN('CA', BTN('CD', BTN('CE'))))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D'))
right = BTN('C', BTN('E', BTN('G', BTN('K'), BTN('L')), BTN('H')), BTN('F', BTN('I'), BTN('J')))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))

left = BTN('B', BTN('D'))
right = BTN('C', BTN('E', BTN('G')), BTN('F'))
tree = BTN('A', left, right)
print(height_largest_complete_subtree(tree))