From 3bba030d53f2ee913fd0687c149a544eb39911b2 Mon Sep 17 00:00:00 2001 From: Matej Focko Date: Thu, 19 May 2022 23:59:50 +0200 Subject: [PATCH] feat: handle root change in rotate function Signed-off-by: Matej Focko --- node.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/node.py b/node.py index 491403f..a1ac28d 100644 --- a/node.py +++ b/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: