feat: handle root change in rotate function
Signed-off-by: Matej Focko <mfocko@redhat.com>
This commit is contained in:
parent
a8de124eeb
commit
3bba030d53
1 changed files with 7 additions and 2 deletions
9
node.py
9
node.py
|
@ -8,6 +8,7 @@ from typing import (
|
|||
Tuple,
|
||||
TypeVar,
|
||||
Protocol,
|
||||
Any,
|
||||
)
|
||||
|
||||
|
||||
|
@ -112,7 +113,7 @@ class Node(Generic[T]):
|
|||
return self
|
||||
|
||||
@staticmethod
|
||||
def rotate_right(x: "Node[T]") -> "Node[T]":
|
||||
def rotate_right(x: "Node[T]", tree: Any = None) -> "Node[T]":
|
||||
parent = x.parent
|
||||
y = x.left
|
||||
# z = x.right
|
||||
|
@ -123,6 +124,8 @@ class Node(Generic[T]):
|
|||
parent.left = y
|
||||
else:
|
||||
parent.right = y
|
||||
elif tree:
|
||||
tree.root = y
|
||||
|
||||
x.left = y.right
|
||||
if x.left:
|
||||
|
@ -135,7 +138,7 @@ class Node(Generic[T]):
|
|||
return y
|
||||
|
||||
@staticmethod
|
||||
def rotate_left(x: "Node[T]") -> "Node[T]":
|
||||
def rotate_left(x: "Node[T]", tree: Any = None) -> "Node[T]":
|
||||
parent = x.parent
|
||||
# y = x.left
|
||||
z = x.right
|
||||
|
@ -146,6 +149,8 @@ class Node(Generic[T]):
|
|||
parent.left = z
|
||||
else:
|
||||
parent.right = z
|
||||
elif tree:
|
||||
tree.root = z
|
||||
|
||||
x.right = z.left
|
||||
if x.right:
|
||||
|
|
Loading…
Reference in a new issue