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 min(self):
  65                 """
  66                 Returns the minimum value of the tree.
  67                 """
  68                 while (self.lnode):
  69                         self = self.lnode
  70                 
  71                 return self.content
  72 
  73         def max(self):
  74                 """
  75                 Returns the maximum value of the tree.
  76                 """
  77                 while (self.rnode):
  78                         self = self.rnode
  79                 
  80                 return self.content
  81 
  82         def search(self, element=None):
  83                 """
  84                 Returns the level of the element on the tree.
  85                 """
  86                 if not element: return -1
  87 
  88                 level = 1
  89 
  90                 while (1):
  91                         if not self: return -1
  92                         if (not self.lnode) and (not self.rnode): return -1
  93 
  94                         if self.content == element: break
  95 
  96                         if self.content <= element: self = self.rnode
  97                         else: self = self.lnode
  98 
  99                         level += 1
 100 
 101                 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))
   8 tree.view()
   9 
  10 print '=' * 25
  11 print "Min: %d" % tree.min()
  12 print "Max: %d" % tree.max()
  13 
  14 sort = randint(1, 100)
  15 print "%d found in %d level." % (sort, tree.search(sort))