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()
class LargerRedBlackBSTMethods(unittest.TestCase):
L = [1, 3, 6, 7, 10, 13, 16]
def setUp(self):
self.st = RedBlackBST()
for x in self.L:
self.st.put(x, x)
def test_put_and_get(self):
st = self.st
for x in st.keys():
self.assertEqual(st.get(x), x)
for x in st.keys():
st.put(x, st.get(x) - 1)
for x in st.keys():
self.assertEqual(st.get(x), x - 1)
self.assertIsNone(st.get(2))
st.put(2, 2)
self.assertEqual(2, st.get(2))
def test_keys(self):
i = 0
for k in self.st.keys():
self.assertEqual(k, self.L[i])
i += 1
def test_min(self):
self.assertEqual(self.st.min(), min(self.L))
def test_max(self):
self.assertEqual(self.st.max(), max(self.L))
def test_delete_min(self):
self.st.delete_min()
self.assertEqual(self.st.min(), min(self.L[1:]))
def test_delete_max(self):
self.st.delete_max()
self.assertEqual(self.st.max(), max(self.L[:-1]))
def test_range(self):
T = [6, 7, 10, 13]
i = 0
for k in self.st.keys_range(5, 13):
self.assertEqual(k, T[i])
i += 1
class HugeRedBlackBSTMethods(unittest.TestCase):
def setUp(self):
random.seed(0)
self.L = random.sample(range(10**6), 10**4)
self.S = sorted(self.L)
self.st = RedBlackBST()
for x in self.L:
self.st.put(x, x)
def test_put_and_get(self):
st = self.st
for x in st.keys():
self.assertEqual(st.get(x), x)
for x in st.keys():
st.put(x, st.get(x) - 1)
for x in st.keys():
self.assertEqual(st.get(x), x - 1)
def test_keys(self):
i = 0
for k in self.st.keys():
self.assertEqual(k, self.S[i])
i += 1
def test_min(self):
self.assertEqual(self.st.min(), min(self.L))
def test_max(self):
self.assertEqual(self.st.max(), max(self.L))
def test_delete_min(self):
self.st.delete_min()
self.assertEqual(self.st.min(), min(self.S[1:]))
def test_delete_max(self):
self.st.delete_max()
self.assertEqual(self.st.max(), max(self.S[:-1]))
def test_min_priority_queue(self):
i = 0
while not self.st.is_empty():
self.assertEqual(self.S[i], self.st.min())
self.st.delete_min()
i += 1
def test_max_priority_queue(self):
i = len(self.S) - 1
while not self.st.is_empty():
self.assertEqual(self.S[i], self.st.max())
self.st.delete_max()
i -= 1
def test_flights(self):
self.st = RedBlackBST()
schedule = [
["09:00:00", "Chicago"],
["09:00:03", "Phoenix"],
["09:00:13", "Houston"],
["09:00:59", "Chicago"],
["09:01:10", "Houston"],
["09:03:13", "Chicago"],
["09:10:11", "Seattle"],
["09:10:25", "Seattle"],
["09:14:25", "Phoenix"],
["09:19:32", "Chicago"],
["09:19:46", "Chicago"],
["09:21:05", "Chicago"],
["09:22:43", "Seattle"],
["09:22:54", "Seattle"],
["09:25:52", "Chicago"],
["09:35:21", "Chicago"],
["09:36:14", "Seattle"],
["09:37:44", "Phoenix"],
]
for time, city in schedule:
self.st.put(time, city)
for time, city in schedule:
self.assertEqual(city, self.st.get(time))
self.assertEqual(len(self.st.keys()), len(schedule))