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
# solving kattis problem 'flights' import sys from datetime import datetime, timedelta from itu.algs4.searching.red_black_bst import RedBlackBST lines = sys.stdin.readlines() n, m = [int(i) for i in lines[0].strip().split()] # read flights into ST (with support for operations) ST = RedBlackBST() for line in lines[1:n + 1]: time, destination = line.strip().split() ST.put(datetime.strptime(time, '%H:%M:%S'), destination) # answer queries for q in lines[n + 1:len(lines)]: query = q.strip().split() command = query[0] # cancel flight by deleting from ST if command == 'cancel': ST.delete(datetime.strptime(query[1], '%H:%M:%S')) # delay a flight by d amount of seconds elif command == 'delay': s, d = query[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)
N = int(sys.argv[2]) if sys.argv[1] == 'Blist': l = blist.sortedlist([i for i in range(10**N)]) elif sys.argv[1] == 'SortedArray': l = SCArray._SortedArray('q', [i for i in range(10**N)]) elif sys.argv[1] == 'SortedList': l = sortedcontainers.SortedList([i for i in range(10**N)]) elif sys.argv[1] == 'AutoLoad': l = SCAutoBalance.SortedList([i for i in range(10**N)]) elif sys.argv[1] == 'BadBisect': l = SCPyBisect.SortedList([i for i in range(10**N)]) elif sys.argv[1] == 'RBSTree': l = RedBlackBST() for i in range(10**N): l.put(i, 1) elif sys.argv[1] == 'SkipList': l = SkipList() for i in range(10**N): l.insert(i, i) else: print("Error, did not supply working thingy") exit(-1) print(f"Testing {sys.argv[1]}") print(len(l)) if sys.argv[3] == 'Base': pass elif sys.argv[3] == 'Add': if sys.argv[1] == 'SkipList': for i in range(1, 10**N): #No duplicates
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 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))