diff --git a/pydatastructs/trees/binary_trees.py b/pydatastructs/trees/binary_trees.py index d10fd4380..d6201009e 100644 --- a/pydatastructs/trees/binary_trees.py +++ b/pydatastructs/trees/binary_trees.py @@ -566,6 +566,9 @@ def _right_rotate(self, j, k): self.tree[self.tree[k].parent].left = k self.tree[j].parent = k self.tree[k].right = j + kp = self.tree[k].parent + if kp is None: + self.root_idx = k def _left_right_rotate(self, j, k): i = self.tree[k].right @@ -617,6 +620,9 @@ def _left_rotate(self, j, k): self.tree[self.tree[k].parent].right = k self.tree[j].parent = k self.tree[k].left = j + kp = self.tree[k].parent + if kp is None: + self.root_idx = k class AVLTree(SelfBalancingBinaryTree): """ @@ -645,9 +651,6 @@ def _right_rotate(self, j, k): super(AVLTree, self)._right_rotate(j, k) self.tree[j].height = max(self.left_height(self.tree[j]), self.right_height(self.tree[j])) + 1 - kp = self.tree[k].parent - if kp is None: - self.root_idx = k if self.is_order_statistic: self.tree[j].size = (self.left_size(self.tree[j]) + self.right_size(self.tree[j]) + 1) @@ -682,9 +685,6 @@ def _left_rotate(self, j, k): self.right_height(self.tree[j])) + 1 self.tree[k].height = max(self.left_height(self.tree[k]), self.right_height(self.tree[k])) + 1 - kp = self.tree[k].parent - if kp is None: - self.root_idx = k if self.is_order_statistic: self.tree[j].size = (self.left_size(self.tree[j]) + self.right_size(self.tree[j]) + 1)