class TestRedBlackBSTMethods(unittest.TestCase): def setUp(self): self.st = RedBlackBST() def test_empty(self): self.assertTrue(self.st.is_empty()) self.st.put("spam", 0) self.assertFalse(self.st.is_empty()) def test_size(self): for i in range(10): self.assertEqual(i, self.st.size()) self.st.put(str(i), i) for i in range(10): self.assertEqual(10, self.st.size()) self.st.put(str(i), i + 1) # key already there: no change in size for i in reversed(range(10)): self.assertEqual(i + 1, self.st.size()) self.st.delete(str(i)) self.assertEqual(0, self.st.size()) def test_rank_select(self): for i in range(0, 2**8 + 2, 2): self.st.put(i, i) self.assertEqual(0, self.st.min()) self.assertEqual(i, self.st.max()) self.assertEqual(i, self.st.select(i // 2)) self.assertEqual(i // 2, self.st.rank(i)) def test_floor_and_ceiling(self): self.st.put(0, 0) self.st.put(2, 2) self.assertEqual(0, self.st.floor(0)) self.assertEqual(0, self.st.floor(1)) self.assertEqual(2, self.st.floor(2)) self.assertEqual(2, self.st.floor(3)) self.assertEqual(0, self.st.ceiling(-1)) self.assertEqual(0, self.st.ceiling(0)) self.assertEqual(2, self.st.ceiling(1)) self.assertEqual(2, self.st.ceiling(2)) def test_exceptions(self): with self.assertRaises(NoSuchElementException): self.st.min() with self.assertRaises(NoSuchElementException): self.st.max() with self.assertRaises(NoSuchElementException): self.st.floor(0) with self.assertRaises(NoSuchElementException): self.st.ceiling(0) self.st.put(0, 0) with self.assertRaises(NoSuchElementException): self.st.floor(-1) with self.assertRaises(NoSuchElementException): self.st.ceiling(1)
class RBSTree(Template): def __init__(self, preload = []): self.li = RedBlackBST() for i in preload: # Init is dumb :( self.add(self, i) def add(self, element): self.li.put(element, 1) def delete(self, index): self.li.delete(self.li.select(index)) def remove(self, element): self.li.delete(element) def rank(self, element): return self.li.rank(element) def select(self, index): return self.li.select(index) def iter(self): return iter(self.li.keys()) def reversed(self): return reversed(self.li.keys()) def count(self, element): return self.li.get(element) def successor(self, value): return self.li.floor(value + 1) def predecessor(self, value): return self.li.ceiling(value - 1) def size(self): return self.li.size()