btree_6.py

#! /usr/bin/python3
# ==================================================================
# create a binary search tree with multiple node types
#
# Notes:
# 1. To simplify the code, all of the nodes have a common key
#    (an integer) for node comparison/insertion. (A sophisticated
#    comparison method could be developed.)
# 2. Nodes are created outside of the binary tree code and inserted
#    into the tree. Thus, The binary tree code is not responsible
#    for the node's correctness.
# 3. Perhaps one way to make the code less confusing to read
#    is for each node have a method to get it's key?
# ==================================================================

class NodeInum:
    '''integer data object'''
    def __init__(self, key, inum):
        self.key  = key
        self.left = None
        self.right = None
        self.inum = inum

class NodeFnum:
    '''float data object'''
    def __init__(self, key, fnum):
        self.key  = key
        self.left = None
        self.right = None
        self.fnum = fnum

class NodeStr:
    '''string data object'''
    def __init__(self, key, str):
        self.key  = key
        self.left = None
        self.right = None
        self.str = str

class NodeUnk:
    '''unknow data object (for debug testing)'''
    def __init__(self, key, data):
        self.key = key
        self.left = None
        self.right = None
        self.data = data

class BinaryTree:
    '''binary tree object'''

    # ---- init object --------------------------

    def __init__(self):
        self.root = None

    # ---- insert object into tree --------------

    def insert(self,newnode):
        if self.root == None:
            self.root = newnode
        else:
            self._insert(newnode,self.root)

    def _insert(self,newnode,curnode):
        if newnode.key < curnode.key:
            if curnode.left == None:
                curnode.left = newnode
            else:
                self._insert(newnode,curnode.left)
        elif newnode.key > curnode.key:
            if curnode.right == None:
                curnode.right = newnode
            else:
                self._insert(newnode,curnode.right)
        else:
            print("Tree node " + str(curnode.key) +
                  " already in the tree")

    # ---- print tree ---------------------------

    def PrintTree(self):
        if self.root == None:
            print("tree empty")
        else:
            self._PrintTree(self.root)

    def _PrintTreeNode(self,node):
        if isinstance(node,NodeInum):
            print("node (key={}) type NodeInum  ({})".
            format(str(node.key),str(node.inum)))
        elif isinstance(node,NodeFnum):
            print("node (key={}) type NodeFnum  ({})".
            format(str(node.key),str(node.fnum)))
        elif isinstance(node,NodeStr):
            print("node (key={}) type NodeStr   ({})".
            format(str(node.key),node.str))
        else:
            print("node (key={}) type unknown   type={}".
            format(str(node.key),type(node).__name__))
            return False
        return True

    def _PrintTree(self,curnode):
        if curnode.left:
            self._PrintTree(curnode.left)
        if not self._PrintTreeNode(curnode):
            pass ##return
        if curnode.right:
            self._PrintTree(curnode.right)

    # ---- print tree structure -----------------

    def PrintTreeStructure(self):
        if self.root == None:
            print("tree empty")
        else:
            self._PrintTreeStructure(self.root,1)

    def _PrintTreeStructureNode(self,node,n):
        if isinstance(node,NodeInum):
            print('({}) {} {}'.format(str(node.key),
            '---'*n,str(node.inum)))
        elif isinstance(node,NodeFnum):
            print('({}) {} {}'.format(str(node.key),
            '---'*n,str(node.fnum)))
        elif isinstance(node,NodeStr):
            print('({}) {} {}'.format(str(node.key),
            '---'*n,node.str))
        else:
            print('({}) {} {}'.format(str(node.key),
            '---'*n,' unknow'))
            return False
        return True

    def _PrintTreeStructure(self,curnode,n):
        if curnode.right:
            self._PrintTreeStructure(curnode.right,n+1)
        if not self._PrintTreeStructureNode(curnode,n):
            pass ##return
        if curnode.left:
            self._PrintTreeStructure(curnode.left,n+1)


# ---- test tree 1
tree = BinaryTree()
tree.insert(NodeInum(6,12))         # root node
tree.insert(NodeFnum(3,3.13159))
tree.insert(NodeStr(5,"abc"))
tree.insert(NodeUnk(4,"unk"))
tree.insert(NodeInum(7,666))
tree.insert(NodeFnum(9,2.71828))
tree.insert(NodeStr(8,"xyz"))
tree.insert(NodeInum(1,"102"))

print("-----------------------------------------------------")

tree.PrintTree()

print("-----------------------------------------------------")

tree.PrintTreeStructure()

print("-----------------------------------------------------")