def test_red_single(self):
tree = TreeSet((2, 1, 5, 3))
tree.discard(3)
self.assertEqual(2, tree.black_height())
self.assertEqual(2, tree.height())
self.assertEqual((1, 2, 5), tuple(tree))
self.assertTrue(tree.is_valid())
def test_simple(self):
tree = TreeSet()
for _ in (2, 1, 3):
tree.add(_)
self.assertEqual(Color.BLACK, tree.root.color)
self.assertEqual(Color.RED, tree.root.left.color)
self.assertEqual(Color.RED, tree.root.right.color)
self.assertEqual(2, tree.height())
self.assertEqual(1, tree.black_height())
self.assertTrue(tree.is_valid())
def test_root_left(self):
tree = TreeSet()
for _ in (5, 3, 6, 1, 4, 2):
tree.add(_)
self.assertEqual(2, tree.black_height())
self.assertEqual(4, tree.height())
tree.discard(5)
self.assertEqual(4, tree.root.key)
self.assertEqual(2, tree.root.left.key)
self.assertEqual(Color.RED, tree.root.left.color)
self.assertEqual(Color.BLACK, tree.root.left.left.color)
self.assertEqual(Color.BLACK, tree.root.left.right.color)
self.assertTrue(tree.is_valid())
def test_black_sibling_red_child(self):
tree = TreeSet()
for _ in (1, 2, 3, 4, 6, 8, 9):
tree.add(_)
self.assertEqual(3, tree.black_height())
self.assertEqual(3, tree.height())
tree.add(5)
tree.add(7)
self.assertEqual(3, tree.black_height())
self.assertEqual(4, tree.height())
tree.discard(9)
self.assertEqual(3, tree.black_height())
self.assertEqual(4, tree.height())
self.assertEqual(4, tree.root.key)
self.assertEqual(6, tree.root.right.key)
self.assertEqual(Color.BLACK, tree.root.right.color)
self.assertEqual(5, tree.root.right.left.key)
self.assertEqual(Color.BLACK, tree.root.right.left.color)
self.assertEqual(8, tree.root.right.right.key)
self.assertEqual(Color.BLACK, tree.root.right.right.color)
self.assertEqual(7, tree.root.right.right.left.key)
self.assertEqual(Color.RED, tree.root.right.right.left.color)
self.assertTrue(tree.is_valid())
def test_all_black(self):
tree = TreeSet()
for _ in (1, 2, 3, 4, 5, 6, 7):
tree.add(_)
self.assertEqual(3, tree.black_height())
self.assertEqual(3, tree.height())
tree.discard(7)
self.assertEqual(2, tree.black_height())
self.assertEqual(4, tree.root.key)
self.assertEqual(Color.RED, tree.root.left.color)
self.assertEqual(Color.BLACK, tree.root.right.color)
self.assertEqual(Color.RED, tree.root.right.left.color)
self.assertEqual(None, tree.root.right.right)
self.assertTrue(tree.is_valid())