finish AVL

shangxunhang · Jul 7, 2017 · 8384cbb · 8384cbb
1 parent d8a0afc
commit 8384cbb
Showing 1 changed file with 56 additions and 26 deletions.
diff --git a/data_structures/AVL/AVL.py b/data_structures/AVL/AVL.py
@@ -39,7 +39,7 @@ def setParent(self, parent):
     def setHeight(self, height):
         self.height = height
 
-    def getHeight(self, height):
+    def getHeight(self):
         return self.height
 
 
@@ -72,38 +72,36 @@ def insert(self, value):
                     if node.getLabel() < dad_node.getLabel():
                         dad_node.setLeft(node)
                         dad_node.setHeight(dad_node.getHeight() + 1)
-
-                        if (dad_node.getRight().getHeight() -
-                                dad_node.getLeft.getHeight() > 1):
-                            self.rebalance(dad_node)
-
                     else:
                         dad_node.setRight(node)
                         dad_node.setHeight(dad_node.getHeight() + 1)
-
-                        if (dad_node.getRight().getHeight() -
-                                dad_node.getLeft.getHeight() > 1):
-                            self.rebalance(dad_node)
                     break
 
     def rebalance(self, node):
-        if (node.getRight().getHeight() -
-                node.getLeft.getHeight() > 1):
-            if (node.getRight().getHeight() >
-                    node.getLeft.getHeight()):
-                pass
-            else:
-                pass
-            pass
-        elif (node.getRight().getHeight() -
-                node.getLeft.getHeight() > 2):
-            if (node.getRight().getHeight() >
-                    node.getLeft.getHeight()):
-                pass
+        height_right = 0
+        height_left = 0
+
+        if node.getRight() is not None:
+            height_right = node.getRight().getHeight()
+
+        if node.getLeft() is not None:
+            height_left = node.getLeft().getHeight()
+
+        if abs(height_left - height_right) > 1:
+            if height_left > height_right:
+                right_child = node.getRight()
+                if (right_child.getLeft().getHeight() >
+                        right_child.getRight().getHeight()):
+                    self.rotate_left(node)
+                else:
+                    self.double_rotate_right(node)
             else:
-                pass
-            pass
-        pass
+                left_child = node.getLeft()
+                if (left_child.getLeft().getHeight() >
+                        left_child.getRight().getHeight()):
+                    self.double_rotate_left(node)
+                else:
+                    self.rotate_right(node)
 
     def rotate_left(self, node):
         # TODO: is this pythonic enought?
@@ -129,3 +127,35 @@ def double_rotate_left(self, node):
     def double_rotate_right(self, node):
         self.rotate_left(node.getLeft().getLeft())
         self.rotate_right(node)
+
+    def empty(self):
+        if self.root is None:
+            return True
+        return False
+
+    def preShow(self, curr_node):
+        if curr_node is not None:
+            self.preShow(curr_node.getLeft())
+            print(curr_node.getLabel(), end=" ")
+            self.preShow(curr_node.getRight())
+
+    def getRoot(self):
+        return self.root
+
+t = AVL()
+t.insert(11)
+t.insert(14)
+t.insert(3)
+t.insert(4)
+t.insert(5)
+t.insert(6)
+t.insert(7)
+t.insert(8)
+t.insert(9)
+t.insert(10)
+t.insert(1)
+t.insert(12)
+t.insert(13)
+t.insert(2)
+t.insert(15)
+t.preShow(t.getRoot())