Pacific-Design.com

    
Home Index

1. Python

2. Binary Tree

Python / Binary Tree /

#!/usr/bin/python

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

    def insert(self, data):
        if data < self.data:
            if self.left is None:
                self.left = Node(data)
            else:
                self.left.insert(data)
        else:
            if self.right is None:
                self.right = Node(data)
            else:
                self.right.insert(data)

    def lookup(self, data, parent=None):
        if data < self.data:
            if self.left is None:
                return None, None
            return self.left.lookup(data, self)
        elif data > self.data:
            if self.right is None:
                return None, None
            return self.right.lookup(data, self)
        else:
            return self, parent

    def children_count(self):
        if node is None:
            return None
        cnt = 0
        if self.left:
            cnt += 1
        if self.right:
            cnt += 1
        return cnt

    def delete(self, data):
        node, parent = self.lookup(data)
        if node is not None:
            children_count = node.children_count()
 
        #1- The node to remove has no child.
        if children_count == 0:
            if parent.left is node:
                parent.left = None
            else:
                parent.right = None

        #2- The node to remove has 1 childi, replace it by child
        elif children_count == 1:
            if node.left:
                n = node.left
            else:
                n = node.right
            if parent.left is node:
                parent.left = n
            else:
                parent.right = n
            del node

        #3- The node to remove has 2 children, find its successor
        else:
            parent = node
            successor = node.right
            while successor.left:
                parent = successor
                successor = successor.left
            # replace node data by its successor data
            node.data = successor.data
            # fix successor's parent's child
            if parent.left == successor:
                parent.left = successor.right
            else:
                parent.right = successor.right

    def print_tree(self):
        if self.left:
            self.left.print_tree()
        print self.data,
        if self.right:
          self.right.print_tree()

if __name__ == "__main__":
    print "Starting to build tree ..."
    root = Node("kevin")
    root.insert("thomas")
    root.insert("duraj")
    root.insert("binary")
    root.insert("search")
    root.insert("tree")
    root.insert("insert")
    root.insert("delete")
    root.insert("lookup")

    node, parent = root.lookup("tree")
    print node.data
    print parent.data
    root.delete("search")
    root.print_tree()