Advertisement
Guest User

Untitled

a guest
Sep 16th, 2019
126
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 5.14 KB | None | 0 0
  1. class Node:
  2. def __init__(self, key, value, balance=0, left=None,right=None):
  3. self.key = key
  4. self.value = value
  5. self.balance = balance
  6. self.left = left
  7. self.right = right
  8.  
  9. class AVL:
  10. def __init__(self):
  11. self.root = None
  12. self.size = 0
  13.  
  14. def put(self, key, value):
  15. self.root, growth = self._put(self.root, key, value)
  16.  
  17. def remove(self, key):
  18. self.root, growth, found = self._remove(self.root, key)
  19. return found
  20.  
  21. def _put(self, pointer, key, value):
  22. if pointer is None:
  23. self.size += 1
  24. return Node(key, value), 1
  25.  
  26. if key < pointer.key:
  27. pointer.left, child_growth = self._put(pointer.left, key, value)
  28. previous_balance = pointer.balance
  29. pointer.balance -= child_growth
  30. growth = 1 if previous_balance <= 0 and pointer.balance < previous_balance else 0 # make it more generic
  31. elif key > pointer.key:
  32. pointer.right, child_growth = self._put(pointer.right, key, value)
  33. previous_balance = pointer.balance
  34. pointer.balance += child_growth
  35. growth = 1 if previous_balance >= 0 and pointer.balance > previous_balance else 0 # make it more generic
  36. else:
  37. pointer.key, pointer.value = key, value
  38.  
  39. rotated_pointer, rotation_growth = self._rotate(pointer)
  40. return rotated_pointer, growth + rotation_growth
  41.  
  42. def _remove(self, pointer, key):
  43. if pointer is None:
  44. return None, 0, False
  45.  
  46. if key < pointer.key:
  47. pointer.left, child_growth, found = self._remove(pointer.left, key)
  48. previous_balance = pointer.balance
  49. pointer.balance -= child_growth
  50. growth = 0 if previous_balance >= 0 or pointer.balance <= previous_balance else -1 # make it more generic
  51. elif key > pointer.key:
  52. pointer.right, child_growth, found = self._remove(pointer.right, key)
  53. previous_balance = pointer.balance
  54. pointer.balance += child_growth
  55. growth = 0 if previous_balance <= 0 or pointer.balance >= previous_balance else -1 # make it more generic
  56. else:
  57. if pointer.left is None and pointer.right is None:
  58. return None, -1, True
  59. elif pointer.left is not None and pointer.right is None:
  60. return child.left, -1, True
  61. elif pointer.left is None and pointer.right is not None:
  62. return child.right, -1, True
  63. else:
  64.  
  65. # def _left_nearest(pointer):
  66. # pointer = pointer.left
  67. # while pointer.right is not None:
  68. # pointer = pointer.right
  69. # return pointer
  70.  
  71. # dummy_key = pointer.right.key
  72. # nearest_pointer = _left_nearest(pointer)
  73.  
  74. def _right_nearest(pointer):
  75. pointer = pointer.right
  76. while pointer.left is not None:
  77. pointer = pointer.left
  78. return pointer
  79.  
  80. dummy_key = pointer.left.key
  81. nearest = _right_nearest(pointer)
  82.  
  83. nearest_key, nearest_value = nearest.key, nearest.value
  84. nearest.key, nearest.value = pointer.key, pointer.value
  85. pointer.key, pointer.value = dummy_key, nearest_value
  86.  
  87. rotated_pointer, growth, found = self._remove(pointer, key)
  88. pointer.key = nearest_key
  89. return rotated_pointer, growth, found
  90.  
  91. rotated_pointer, rotation_growth = self._rotate(pointer)
  92. return rotated_pointer, growth + rotation_growth, True
  93.  
  94. def _rotate(self, pointer):
  95. rotated_pointer = pointer
  96. rotation_growth = 0
  97. if pointer.balance <= -2:
  98. if pointer.left.balance > 0:
  99. rotated_pointer, growth = self._left(rotated_pointer.left)
  100. rotation_growth += growth
  101. rotated_pointer, growth = self._right(rotated_pointer)
  102. rotation_growth += growth
  103. elif pointer.balance >= 2:
  104. if pointer.right.balance < 0:
  105. rotated_pointer, growth = self._right(rotated_pointer.right)
  106. rotation_growth += growth
  107. rotated_pointer, growth = self._left(rotated_pointer)
  108. rotation_growth += growth
  109.  
  110. return rotated_pointer, rotation_growth
  111.  
  112. def _left(self, pointer):
  113. child = pointer.right
  114. pointer.right = child.left
  115. child.left = pointer
  116. growth = -1 if pointer.balance >= 2 else 0 if pointer.balance == 1 else 1
  117. pointer.balance = pointer.balance - 1 - max(child.balance, 0)
  118. child.balance = child.balance - 1 + min(pointer.balance, 0)
  119. return child, growth
  120.  
  121. def _right(self, pointer):
  122. child = pointer.left
  123. pointer.left = child.right
  124. child.right = pointer
  125. growth = -1 if pointer.balance >= -2 else 0 if pointer.balance == -1 else 1
  126. pointer.balance = pointer.balance + 1 - min(child.balance, 0)
  127. child.balance = child.balance + 1 + max(pointer.balance, 0)
  128. return child, growth
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement