feat: handle root change in rotate function

Signed-off-by: Matej Focko <mfocko@redhat.com>
This commit is contained in:
Matej Focko 2022-05-19 23:59:50 +02:00
parent a8de124eeb
commit 3bba030d53
Signed by: mfocko
GPG key ID: 7C47D46246790496

View file

@ -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: