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 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)
time = datetime.strptime(s, '%H:%M:%S')
val = ST.get(time)
ST.delete(time)
updated_t = time + timedelta(seconds=int(d))
ST.put(updated_t, val)
# reroute a flight such that the destination changes for a given key
elif command == 'reroute':
s, c = query[1:]
ST.put(datetime.strptime(s, '%H:%M:%S'), c)
# query ST for time, print destination if time exists, else '-'
elif command == 'destination':
time = datetime.strptime(query[1], '%H:%M:%S')
if ST.get(time) is not None:
print(ST.get(time))
else:
print('-')
# query ST for next available flight, which is equal to computing the ceiling of the given key
elif command == 'next':
time = datetime.strptime(query[1], '%H:%M:%S')
ceiling = ST.ceiling(time)
print(ceiling.time(), ST.get(ceiling))
# query ST for number of flights in given time interval
elif command == 'count':
t1, t2 = datetime.strptime(query[1], '%H:%M:%S'), datetime.strptime(
query[2], '%H:%M:%S')
print(ST.size_range(t1, t2))