1 #!/usr/bin/python
   2 #
   3 # BSTree.py
   4 
   5 import random
   6 
   7 class BSTree:
   8         """
   9         Author: Danillo Souza
  10         Email: danillo012@gmail.com
  11 
  12         DESCRIPTION
  13                 Class that implements Binary Search Tree functionality.
  14 
  15         TODO
  16                 * improve the 'view' method to be more legible and nice;
  17         """
  18         def __init__(self, content=None):
  19                 """
  20                 Initialize the tree with a valid value(can't be None).
  21                 """
  22                 if not content: self.content = 0
  23                 else: self.content = content
  24                 
  25                 self.lnode = None
  26                 self.rnode = None
  27 
  28         def add(self, content=None):
  29                 """
  30                 Add an element to the tree. In the right if the value if is bigger than the
  31                 actual, and in the left if its less than the actual.
  32                 """
  33                 if not content: return -1
  34 
  35                 if content > self.content:
  36                         if self.rnode: self.rnode.add(content)
  37                         else: self.rnode = BSTree(content)
  38 
  39                 if content <= self.content:
  40                         if self.lnode: self.lnode.add(content)
  41                         else: self.lnode = BSTree(content)
  42 
  43         def view(self, level=1):
  44                 """
  45                 Print the tree on the screen.
  46                 """
  47                 arrow = "---" * level
  48                 print "|%s>%s" % (arrow, self.content)
  49 
  50                 if self.rnode: self.rnode.view(level+1)
  51                 if self.lnode: self.lnode.view(level+1)
  52 
  53                 return
  54 
  55         def size(self):
  56                 """
  57                 Returns the size of the tree.
  58                 """
  59                 tsize = 1
  60 
  61                 if self.rnode: tsize += self.rnode.size()
  62                 if self.lnode: tsize += self.lnode.size()
  63 
  64                 return tsize
  65 
  66         def search(self, element=None):
  67                 """
  68                 Returns the level of the element on the tree.
  69                 """
  70                 if not element: return -1
  71 
  72                 level = 1
  73 
  74                 while (1):
  75                         if not self: return -1
  76                         if (not self.lnode) and (not self.rnode): return -1
  77 
  78                         if self.content == element: break
  79 
  80                         if self.content <= element: self = self.rnode
  81                         else: self = self.lnode
  82 
  83                         level += 1
  84 
  85                 return level

   1 #!/usr/bin/python
   2 import BSTree
   3 from random import randint
   4 
   5 if __name__ == '__main__':
   6         tree = BSTree(50)
   7         for i in range(100):
   8                 tree.add(randint(1, 100))0
   9         tree.view()
  10         
  11         print '=' * 25
  12         sort = randint(1, 100)
  13         print "%d found in %d level." % (sort, tree.search(sort))