def test_basic(self): int_tree = interval_tree_create() interval_insert(int_tree, Interval(0, 3)) interval_insert(int_tree, Interval(5, 8)) interval_insert(int_tree, Interval(6, 10)) interval_insert(int_tree, Interval(8, 9)) interval_insert(int_tree, Interval(15, 23)) interval_insert(int_tree, Interval(16, 21)) interval_insert(int_tree, Interval(17, 19)) interval_insert(int_tree, Interval(19, 20)) interval_insert(int_tree, Interval(25, 30)) interval_insert(int_tree, Interval(26, 26)) self.assertEqual( Interval(6, 10), interval_search_exactly(int_tree, Interval(6, 10)).data) self.assertEqual(int_tree.nil, interval_search_exactly(int_tree, Interval(10, 12))) node = interval_search(int_tree, Interval(10, 12)) self.assertEqual(Interval(6, 10), node.data) node = interval_search_min(int_tree, Interval(10, 12)) self.assertEqual(Interval(6, 10), node.data) node = interval_search(int_tree, Interval(9, 17)) self.assertTrue(node.data in (Interval(6, 10), Interval(8, 9), Interval(15, 23), Interval(16, 21), Interval(17, 19))) node = interval_search_min(int_tree, Interval(9, 17)) self.assertEqual(Interval(6, 10), node.data) node = interval_search_min(int_tree, Interval(17, 19)) self.assertEqual(Interval(15, 23), node.data) node = interval_search(int_tree, Interval(12, 14)) self.assertEqual(int_tree.nil, node) node = interval_search_min(int_tree, Interval(12, 14)) self.assertEqual(int_tree.nil, node) node = interval_search(int_tree, Interval(31, 100)) self.assertEqual(int_tree.nil, node) node = interval_search_min(int_tree, Interval(31, 100)) self.assertEqual(int_tree.nil, node) interval_pop( int_tree, rb_search(int_tree, interval_default_key(Interval(15, 23)))) node = interval_search_min(int_tree, Interval(17, 19)) self.assertEqual(Interval(16, 21), node.data) interval_pop(int_tree, rb_search(int_tree, interval_default_key(Interval(0, 3)))) node = interval_search(int_tree, Interval(0, 3)) self.assertEqual(int_tree.nil, node) node = interval_search_min(int_tree, Interval(0, 3)) self.assertEqual(int_tree.nil, node)
def __setitem__(self, k, v): already = rb_search(self._rbt, k) if already != self._rbt.nil: already.data.v = v else: rb_insert(self._rbt, _KeyValuePair(k, v)) self._len += 1
def pop(self, k): rbt = self._rbt node = rb_search(rbt, k) if node == rbt.nil: raise KeyError(str(k)) rb_pop(rbt, node) self._len -= 1
def __getitem__(self, k): rbt = self._rbt node = rb_search(self._rbt, k) if node == rbt.nil: raise KeyError(str(k)) return node.data.v
def __contains__(self, k): rbt = self._rbt return rb_search(rbt, k) != rbt.nil
def interval_search_exactly(int_tree: RBTree, interval: Interval): """ Ex 14.3-5. Search for an interval in the given interval tree that's equal to the given interval. """ return rb_search(int_tree, interval_default_key(interval))