Advertisement
frolkin28

AVL tree

May 21st, 2020 (edited)
126
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 5.24 KB | None | 0 0
  1. class Node:
  2.     def __init__(self, value, left=None, right=None):
  3.         self.value = value
  4.         self.left = left
  5.         self.right = right
  6.         self.hight = 0
  7.  
  8.  
  9. class Tree:
  10.     def __init__(self):
  11.         self.root = None
  12.         self.size = 0
  13.  
  14.     def add(self, node, value):
  15.         if node is None:
  16.             return Node(value)
  17.         else:
  18.             if value < node.value:
  19.                 node.left = self.add(node.left, value)
  20.             else:
  21.                 node.right = self.add(node.right, value)
  22.         return self.balance(node)
  23.  
  24.     def build_tree(self, path):
  25.         with open(path, 'r') as f:
  26.             node_list = f.read().strip()
  27.  
  28.         for value in [int(i) for i in node_list.split()]:
  29.             self.root = self.add(self.root, value)
  30.  
  31.     def _print_tree(self, node, level, p_level):
  32.         printed = False
  33.  
  34.         if p_level == level:
  35.             print(f'{node.value}', end=' ')
  36.             printed = True
  37.         else:
  38.             if node.left:
  39.                 printed = self._print_tree(node.left, level+1, p_level)
  40.             if node.right:
  41.                 printed = (self._print_tree(
  42.                     node.right, level+1, p_level)) or printed
  43.  
  44.         return printed
  45.  
  46.     def print_tree(self, node=None):
  47.         if node is None:
  48.             level = 0
  49.             while (self._print_tree(self.root, 0, level)):
  50.                 level += 1
  51.                 print()
  52.         else:
  53.             level = 0
  54.             while (self._print_tree(node, 0, level)):
  55.                 level += 1
  56.                 print()
  57.  
  58.     def update_height(self, node=None):
  59.         if node == None:
  60.             return 0
  61.         else:
  62.             node.height = max(self.update_height(node.left),
  63.                               self.update_height(node.right)) + 1
  64.             return node.height
  65.  
  66.     def height(self, node):
  67.         if node == None:
  68.             return 0
  69.         else:
  70.             return node.height
  71.  
  72.     def bfactor(self, node):
  73.         if not node:
  74.             return 0
  75.         return self.height(node.right) - self.height(node.left)
  76.  
  77.     def right_rotate(self, node):
  78.         if node.left is None:
  79.             raise ValueError
  80.         temp = node.left
  81.         node.left = temp.right
  82.         temp.right = node
  83.         self.update_height(self.root)
  84.         return temp
  85.  
  86.     def left_rotate(self, node):
  87.         if node.right is None:
  88.             raise ValueError
  89.         temp = node.right
  90.         node.right = temp.left
  91.         temp.left = node
  92.         self.update_height(self.root)
  93.         return temp
  94.  
  95.     def double_rotate_right(self, node):
  96.         node.right = self.right_rotate(node.right)
  97.         return self.left_rotate(node)
  98.  
  99.     def double_rotate_left(self, node):
  100.         node.left = self.left_rotate(node.left)
  101.         return self.right_rotate(node)
  102.  
  103.     def balance(self, node):
  104.         self.update_height(self.root)
  105.         if self.bfactor(node) == 2:
  106.             if self.bfactor(node.right) < 0:
  107.                 node.right = self.right_rotate(node.right)
  108.             return self.left_rotate(node)
  109.         elif self.bfactor(node) == -2:
  110.             if self.bfactor(node.left) > 0:
  111.                 node.left = self.left_rotate(node.left)
  112.             return self.right_rotate(node)
  113.  
  114.         return node
  115.  
  116.     def search(self, value, node):
  117.         res = None
  118.         if node.value == value:
  119.             return node
  120.         if node.right:
  121.             res = self.search(value, node.right)
  122.         if node.left:
  123.             temp = self.search(value, node.left)
  124.             if temp is not None:
  125.                 res = temp
  126.         return res
  127.  
  128.     def delete(self, root, value):
  129.  
  130.         if not root:
  131.             return root
  132.  
  133.         elif value < root.value:
  134.             root.left = self.delete(root.left, value)
  135.  
  136.         elif value > root.value:
  137.             root.right = self.delete(root.right, value)
  138.  
  139.         else:
  140.             if root.left is None:
  141.                 temp = root.right
  142.                 root = None
  143.                 return temp
  144.  
  145.             elif root.right is None:
  146.                 temp = root.left
  147.                 root = None
  148.                 return temp
  149.  
  150.             temp = self.getMinValueNode(root.right)
  151.             root.value = temp.value
  152.             root.right = self.delete(root.right,
  153.                                      temp.value)
  154.  
  155.         if root is None:
  156.             return root
  157.  
  158.         root.height = 1 + max(self.height(root.left),
  159.                               self.height(root.right))
  160.  
  161.         balance = self.bfactor(root)
  162.  
  163.         if balance > 1 and self.bfactor(root.left) >= 0:
  164.             return self.right_rotate(root)
  165.  
  166.         if balance < -1 and self.bfactor(root.right) <= 0:
  167.             return self.left_rotate(root)
  168.  
  169.         if balance > 1 and self.bfactor(root.left) < 0:
  170.             root.left = self.left_rotate(root.left)
  171.             return self.right_rotate(root)
  172.  
  173.         if balance < -1 and self.bfactor(root.right) > 0:
  174.             root.right = self.right_rotate(root.right)
  175.             return self.left_rotate(root)
  176.  
  177.         return root
  178.  
  179.     def getMinValueNode(self, root):
  180.         if root is None or root.left is None:
  181.             return root
  182.  
  183.         return self.getMinValueNode(root.left)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement