def test_right_rotation_with_two_nodes(self):
"""
Right rotation
Only two nodes
Successfully
N N
| |
(10) (5)
/ \ / \
/ \ / \
(5) N -> N (10)
/ \ / \
/ \ / \
N N N N
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(10)
red_black_tree.insert(5)
node_to_rotate = red_black_tree.search(10)
red_black_tree._right_rotation(node_to_rotate)
new_left_node = red_black_tree.search(5)
new_right_node = red_black_tree.search(10)
assert new_right_node.parent is new_left_node
assert new_right_node.right is None
assert new_right_node.left is None
assert new_left_node.parent is None
assert new_left_node.right is new_right_node
assert new_left_node.left is None
def test_cases_one_through_three():
rbt = RedBlackTree()
# Case 1, root is head
rbt.insert(10)
assert rbt.color is BLACK
# Case 2, current node's parent is black
rbt.insert(12)
assert rbt.color is BLACK
assert rbt.rightchild.color is RED
# Case 3, parent and uncle are both RED
rbt = RedBlackTree()
rbt.insert(10)
rbt.insert(8)
rbt.insert(12)
rbt.insert(11)
# . 10(B)
# . / \
# . 8(B) 12(B)
# . /
# . 11(R)
child = rbt.rightchild.leftchild
assert child.color is RED
assert rbt.rightchild is child.parent
assert child.parent.color is BLACK
assert child.uncle is rbt.leftchild
assert child.uncle.color is BLACK
assert rbt.color is BLACK
def test_insertion():
rbt = RedBlackTree()
for i in xrange(256):
rbt.insert(random.randint(0, 1e6))
assert rbt.size() == 256
assert abs(rbt.balance()) <= 1
assert rbt.depth() <= 10
check_validity(rbt)
def test_deletion():
numberpool = [random.randint(0, 1e4) for i in xrange(256)]
rbt = RedBlackTree()
for i in numberpool:
rbt.insert(i)
check_validity(rbt)
random.shuffle(numberpool)
for i in numberpool:
rbt.delete(i)
check_validity(rbt)
def test_insert_in_empty_tree(self):
"""
Insert in empty tree
Successfully
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(5)
assert red_black_tree._root.val == 5
assert red_black_tree._root.color == Color.BLACK
assert red_black_tree._root.left is None
assert red_black_tree._root.right is None
def test_insert_in_left_subtree(self):
"""
Insert in left subtree
Successfully
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(5)
red_black_tree.insert(4)
left_inserted_node = red_black_tree.search(4)
right_tree = red_black_tree.search(5)
assert left_inserted_node.parent is right_tree
assert right_tree.left is left_inserted_node
def test_sibling():
rbt = RedBlackTree()
rbt.insert(10) # . 10
rbt.insert(12) # . / \
rbt.insert(8) # . 8 12
rbt.insert(13) # . / / \
rbt.insert(7) # . 7 11 13
rbt.insert(11)
assert rbt.sibling is None
assert rbt.leftchild.sibling.value == 12
assert rbt.rightchild.sibling.value == 8
assert rbt.leftchild.leftchild.sibling is None
assert rbt.rightchild.rightchild.sibling.value == 11
def test_find_func_search_for_not_existing(self):
"""
Searching for not existing value
Successfully
"""
num_of_elements = 100
possible_values = list(range(-100000, 100000))
elements = [
random.choice(possible_values) for _ in range(num_of_elements)
]
different_element = 100001
red_black_tree = RedBlackTree()
for elem in elements:
red_black_tree.insert(elem)
assert red_black_tree.search(different_element) is None
def test_find_func_all_found(self):
"""
Searching all inserted values
Successfully
"""
num_of_elements = 100
possible_values = list(range(-100000, 100000))
elements = [
random.choice(possible_values) for _ in range(num_of_elements)
]
red_black_tree = RedBlackTree()
for elem in elements:
red_black_tree.insert(elem)
random.shuffle(elements)
for elem in elements:
assert red_black_tree.search(elem).val == elem
def test_cases_four_and_five():
rbt = RedBlackTree()
rbt.insert(10)
rbt.insert(12)
rbt.insert(11)
# Origin Rotate for case 4 Rotate again Case 5
# 10(B) | 10(B) | 11(B)
# . \ | \ | / \
# . 12(R) | 11(R) | 10(R) 12(R)
# . / | \ |
# 11(R) | 12(R) |
assert rbt.color is BLACK
assert rbt.value == 11
assert rbt.leftchild.color is RED
assert rbt.leftchild.value == 10
assert rbt.rightchild.color is RED
assert rbt.rightchild.value == 12
def test_right_rotation_with_one_node(self):
"""
Right rotation
Only one node
Successfully
N N
| |
(5) -> (5)
/ \ / \
/ \ / \
(N) N (N) N
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(5)
node_to_rotate = red_black_tree.search(5)
red_black_tree._right_rotation(node_to_rotate)
new_node = red_black_tree.search(5)
assert new_node is node_to_rotate
def count_time():
"""
Count time of insertion and searching
Plot graphics
:return: void
"""
repetitions = 1000
num_of_max_elements = 500
min_possible_value = -100000
max_possible_value = 100000
possible_values = list(range(min_possible_value, max_possible_value))
number_of_elements = [num for num in range(num_of_max_elements)]
all_times_insert = [0] * num_of_max_elements
all_times_search = [0] * num_of_max_elements
for _ in range(repetitions):
red_black_tree = RedBlackTree()
for cur_number_of_elements in number_of_elements:
start_time_insertion = time.perf_counter()
red_black_tree.insert(random.choice(possible_values))
end_time_insertion = time.perf_counter()
all_times_insert[cur_number_of_elements] += end_time_insertion - start_time_insertion
start_time_search = time.perf_counter()
red_black_tree.search(random.choice(possible_values))
end_time_search = time.perf_counter()
all_times_search[cur_number_of_elements] += end_time_search - start_time_search
all_times_insert = [time_insert / repetitions for time_insert in all_times_insert]
all_times_search = [time_search / repetitions for time_search in all_times_search]
plt.plot(number_of_elements, all_times_insert)
plt.plot(number_of_elements, all_times_search)
plt.xlabel("Num of elements")
plt.ylabel("Time")
plt.legend(["Insertion", "Search"])
plt.savefig("time.png")
plt.show()
def test_p_right_child_of_g_and_k_left_child_of_p(self):
"""
Insert
Parent is right child of grandParent ant inserted node
is left child of parent
Successfully
61-B 61-B
/ \ / \
/ \ / \
52-B 85-B -> + 87 -> 52-B 87-B
\ / \
\ / \
93-R 85-R 93-R
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(61)
red_black_tree.insert(52)
red_black_tree.insert(85)
red_black_tree.insert(93)
red_black_tree.insert(87)
node_61 = red_black_tree.search(61)
node_52 = red_black_tree.search(52)
node_85 = red_black_tree.search(85)
node_93 = red_black_tree.search(93)
node_87 = red_black_tree.search(87)
assert node_93.color == Color.RED
assert node_85.color == Color.RED
assert node_87.color == Color.BLACK
assert node_52.color == Color.BLACK
assert node_61.color == Color.BLACK
def test_left_rotation(self):
"""
Left rotation
All nodes(Common case)
Successfully
N N
| |
(5) (7)
/ \ / \
/ \ / \
4 (7) -> (5) 8
/ \ / \
/ \ / \
6 8 4 6
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(5)
red_black_tree.insert(4)
red_black_tree.insert(7)
red_black_tree.insert(6)
red_black_tree.insert(8)
node_to_rotate = red_black_tree.search(5)
red_black_tree._left_rotation(node_to_rotate)
new_left_node = red_black_tree.search(5)
new_right_node = red_black_tree.search(7)
assert new_right_node.parent is None
assert new_right_node.right.val == 8
assert new_right_node.right.parent is new_right_node
assert new_right_node.left is new_left_node
assert new_left_node.parent is new_right_node
assert new_left_node.left.val == 4
assert new_left_node.left.parent is new_left_node
assert new_left_node.right.val == 6
assert new_left_node.right.parent is new_left_node
def test_right_rotation(self):
"""
Right rotation
All nodes(Common case)
Successfully
N N
| |
(8) (6)
/ \ / \
/ \ / \
(6) 9 -> 4 (8)
/ \ / \
/ \ / \
4 7 7 9
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(8)
red_black_tree.insert(9)
red_black_tree.insert(6)
red_black_tree.insert(4)
red_black_tree.insert(7)
node_to_rotate = red_black_tree.search(8)
red_black_tree._right_rotation(node_to_rotate)
new_left_node = red_black_tree.search(6)
new_right_node = red_black_tree.search(8)
assert new_right_node.parent is new_left_node
assert new_right_node.right.val == 9
assert new_right_node.right.parent is new_right_node
assert new_right_node.left.val == 7
assert new_right_node.left.parent is new_right_node
assert new_left_node.left.val == 4
assert new_left_node.left.parent is new_left_node
assert new_left_node.right is new_right_node
def create_tree(self):
red_black_tree = RedBlackTree()
red_black_tree.insert(14)
red_black_tree.insert(35)
red_black_tree.insert(10)
red_black_tree.insert(19)
red_black_tree.insert(31)
red_black_tree.insert(42)
return red_black_tree
def test_grandparent():
rbt = RedBlackTree()
assert rbt.parent is None
assert rbt.grandparent is None
rbt.insert(10)
rbt.insert(11)
rbt.insert(9)
rbt.insert(12)
assert rbt.rightchild.grandparent is None
assert rbt.rightchild.rightchild.grandparent is rbt
def test_insert_already_existing_value(self):
"""
Insert already existing value
Successfully
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(5)
red_black_tree.insert(6)
red_black_tree.insert(7)
red_black_tree.insert(8)
red_black_tree.insert(9)
searching_elem = 9
red_black_tree.insert(searching_elem)
num_of_necessary_nodes = 0
list_of_nodes = red_black_tree.travers_tree("LNR")
for val in list_of_nodes:
if val == searching_elem:
num_of_necessary_nodes += 1
assert num_of_necessary_nodes == 1
def test_preliminary_insert():
rbt = RedBlackTree()
assert rbt.color == BLACK
rbt.insert(10)
assert rbt.value == 10
assert rbt.color == BLACK # Head stays black
rbt.insert(8)
assert rbt.leftchild.value == 8
assert rbt.leftchild.color == RED
rbt.insert(12)
assert rbt.rightchild.value == 12
assert rbt.rightchild.color == RED
def test_insert_p_red_u_red(self):
"""
Insert
Parent is red and Uncle is red
Successfully
61-B 61-B
/ \ / \
/ \ / \
52-B 85-B -> + 100 -> 52-B 85-R
/ \ / \
/ \ / \
76-R 93-R 76-B 93-B
\
\
100-R
"""
red_black_tree = RedBlackTree()
red_black_tree.insert(61)
red_black_tree.insert(52)
red_black_tree.insert(85)
red_black_tree.insert(76)
red_black_tree.insert(93)
red_black_tree.insert(100)
node_61 = red_black_tree.search(61)
node_52 = red_black_tree.search(52)
node_85 = red_black_tree.search(85)
node_76 = red_black_tree.search(76)
node_93 = red_black_tree.search(93)
node_100 = red_black_tree.search(100)
assert node_100.color == Color.RED
assert node_93.color == Color.BLACK
assert node_76.color == Color.BLACK
assert node_85.color == Color.RED
assert node_52.color == Color.BLACK
assert node_61.color == Color.BLACK
def test_uncle():
rbt = RedBlackTree()
assert rbt.uncle is None
rbt.insert(10) # . 10
rbt.insert(12) # . / \
rbt.insert(8) # . 8 12
rbt.insert(13) # . / \ / \
rbt.insert(7) # . 7 9 11 13
rbt.insert(11)
rbt.insert(9)
assert rbt.rightchild.uncle is None
assert rbt.leftchild.uncle is None
assert rbt.rightchild.rightchild.uncle.value == 8
assert rbt.rightchild.leftchild.uncle.value == 8
assert rbt.leftchild.rightchild.uncle.value == 12
assert rbt.leftchild.leftchild.uncle.value == 12
def insert_tree(data):
root = Node(data[0], BLACK)
rbt = RedBlackTree(root)
for i in data[1:]:
rbt.insert(i)
return rbt
def insert_tree(data):
root = Node(data[0], BLACK)
rbt = RedBlackTree(root)
for i in data[1:]:
rbt.insert(i)
return rbt
current = None
while True:
clock.tick(30)
screen.fill((255, 255, 255))
for i in pygame.event.get():
if i.type == pygame.QUIT:
sys.exit()
if i.type == pygame.KEYDOWN:
if i.key == pygame.K_ESCAPE:
sys.exit()
if i.key == pygame.K_i:
print("Enter key(must be number): ")
key = float(input())
print("Enter value: ")
value = input()
T.insert(Node(key, value))
if i.key == pygame.K_f:
if current != None:
current.looking = False
print("Enter key(must be number): ")
key = float(input())
current = T.find(key)
if current != None:
current.looking = True
if i.key == pygame.K_r:
print("Enter key(must be number): ")
T.remove(float(input()))
if i.key == pygame.K_p:
T.printInOrder()
drawBT(T.root)
pygame.display.flip()
