Binary Tree Using Classes

An object-oriented approach

In [7]:
class BinaryTree:
    def __init__(self, key):
        self.key = key
        self.left_child = None
        self.right_child = None
        
    def insert_left(self, key):
        if self.left_child is None:
            self.left_child = BinaryTree(key)
        else:  # if there IS a left child
            t = BinaryTree(key)
            t.left_child = self.left_child
            self.left_child = t
            
    def insert_right(self, key):
        if self.right_child is None:
            self.right_child = BinaryTree(key)
        else:  # if there IS a right child
            t = BinaryTree(key)
            t.right_child = self.right_child
            self.right_child = t
            
    def get_right_child(self):
        return self.right_child
    
    def get_left_child(self):
        return self.left_child
    
    def get_root_val(self):
        return self.key
    
    def set_root_val(self, new_key):
        self.key = new_key
    
    def __repr__(self):
        return f"BinaryTree({self.key!r})"

Testing

In [8]:
r = BinaryTree("a")
In [9]:
r
Out[9]:
BinaryTree('a')
In [12]:
print("root value:", repr(r.get_root_val()))
print("left child:", r.get_left_child())
print("right child:", r.get_right_child())
root value: 'a'
left child: None
right child: None
In [13]:
r.insert_left('b')
print("left child:", r.get_left_child())
left child: BinaryTree('b')
In [14]:
r.insert_right('c')
print("right child:", r.get_right_child())
right child: BinaryTree('c')
In [15]:
r
Out[15]:
BinaryTree('a')
In [16]:
r.get_right_child()
Out[16]:
BinaryTree('c')

Visualizing with graphviz

In [28]:
from graphviz import Digraph
from pythonds9.basic.stack import Stack
In [32]:
def viz_tree(r):
    stack = Stack()
    g = Digraph(node_attr={'shape': 'record', 'height': '.1'})
    _id = 0
    current_node = r
    leftward = True
    current_root_num = 0
    
    while True:
        if current_node:
            stack.push((_id, current_node))

            #g.node(f'node{_id}', 
            #       f'<f0>|<f1> {current_node.key}:{current_node.payload} ({current_node.balance_factor})|<f2> ')
             
            g.node(f'node{_id}', 
                   f'<f0>|<f1> {current_node.key}|<f2> ')

        
            if _id >= 1:
                g.edge(f'node{current_root_num}:f{0 if leftward else 2}',
                       f'node{_id}:f1')
                
            leftward = True
            current_node = current_node.left_child  # left
            current_root_num = _id
            _id += 1

        if current_node is None and not stack.is_empty():
            count, popped_node = stack.pop()
            if popped_node.right_child:
                current_root_num = count
                current_node = popped_node.right_child  # right
                leftward = False
            
        if current_node is None and stack.is_empty():
            break      

    return g
In [33]:
r
Out[33]:
BinaryTree('a')
In [34]:
g = viz_tree(r)
In [35]:
g
Out[35]:
node0 a node1 b node0:f0->node1:f1 node2 c node0:f2->node2:f1
In [36]:
r.get_left_child().insert_right("z")
In [37]:
g = viz_tree(r)
In [38]:
g
Out[38]:
node0 a node1 b node0:f0->node1:f1 node3 c node0:f2->node3:f1 node2 z node1:f2->node2:f1

Testing viz_tree from pythonds9

In [39]:
from pythonds9.tree.utils import viz_tree as viz_tree_obj
In [40]:
g2 = viz_tree_obj(r)
In [41]:
g2
Out[41]:
node0 a node1 b node0:f0->node1:f1 node3 c node0:f2->node3:f1 node2 z node1:f2->node2:f1