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

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