def test_inserting(self):
node = TreapNode(10, "10")
node = node.get_with_inserted(20, "20")
node = node.get_with_inserted(102, "102")
node = node.get_with_inserted(101, "101")
node = node.get_with_inserted(2, "2")
return self.assertTrue(2 in node)
def test_get_no_key_in_treap(self):
node = TreapNode(10, "100")
with self.assertRaises(KeyError) as context:
node.get(20)
self.assertTrue(
check_error_message(context, "Key '20' isn't in the treap."))
def test_update_no_key_in_treap(self):
node = TreapNode(10, "100")
with self.assertRaises(KeyError) as context:
node.update(20, "100")
self.assertTrue(
check_error_message(
context,
"Unable to update because there is no key '20' in the treap"))
def __setitem__(self, key: K, value: V) -> None:
if self._top is not None and key in self:
self._top.update(key, value)
else:
if self._top is None:
self._top = TreapNode(key, value)
else:
self._top = self._top.get_with_inserted(key, value)
self._length += 1
def insert(self, key: K, value: V, priority: P) -> None:
"""Inserts a new pair (key, value) with a priority."""
if key in self:
raise KeyError(f"The key '{key}' is already in the tree")
if self._top is None:
self._top = TreapNode(key, value)
else:
self._top = self._top.get_with_inserted(key, value, priority)
self._length += 1
class Treap(MutableMapping[K, V], Generic[K, P]):
_top: Optional["TreapNode[K, P]"]
_length: int
def __init__(self):
self._length = 0
self._top = None
def __setitem__(self, key: K, value: V) -> None:
if self._top is not None and key in self:
self._top.update(key, value)
else:
if self._top is None:
self._top = TreapNode(key, value)
else:
self._top = self._top.get_with_inserted(key, value)
self._length += 1
def insert(self, key: K, value: V, priority: P) -> None:
"""Inserts a new pair (key, value) with a priority."""
if key in self:
raise KeyError(f"The key '{key}' is already in the tree")
if self._top is None:
self._top = TreapNode(key, value)
else:
self._top = self._top.get_with_inserted(key, value, priority)
self._length += 1
def __delitem__(self, key: K) -> None:
if self._top is None:
raise KeyError("Treap is empty.")
if key not in self:
raise KeyError(f"Key '{key}' isn't in the treap.")
self._top = self._top.get_with_removed(key)
self._length -= 1
def __getitem__(self, key: K) -> V:
if self._top is None:
raise KeyError("Treap is empty.")
if key not in self:
raise KeyError(f"Key '{key}' isn't in the treap.")
return self._top.get(key)
def __len__(self) -> int:
return self._length
def __iter__(self) -> Iterator[K]:
"""Returns an iterator that traverses the keys in ascending order"""
if self._top is None:
yield from ()
else:
yield from self._top
def __reversed__(self) -> Iterator[K]:
"""Returns an iterator that traverses the keys in descending order"""
if self._top is None:
yield from ()
else:
yield from reversed(self._top)
# mypy requires the argument to be of type "object" ("Any") to match the "Mapping" supertype
def __contains__(self, key: Any) -> bool:
return self._top is not None and cast("K", key) in self._top
def test_contains_after_removing(self):
node = TreapNode(5, "20")
node = node.get_with_inserted(20, "20")
node = node.get_with_removed(5)
return self.assertFalse(5 in node)
def test_get_one_node(self):
node = TreapNode(10, "10")
return self.assertEqual(node.get(10), "10")
def test_update(self):
node = TreapNode(10, "10")
node = node.get_with_inserted(20, "20")
node = node.get_with_inserted(-10, "-10")
node.update(20, 100)
return self.assertEqual(node.get(20), 100)
def test_removing(self):
node = TreapNode(10, "10")
node = node.get_with_inserted(20, "20")
node = node.get_with_removed(10)
return self.assertFalse(10 in node)
def test_contains_several_nodes(self):
node = TreapNode(10, "10")
node = node.get_with_inserted(20, "30")
node = node.get_with_inserted(100, "30")
return self.assertTrue(100 in node)
def test_contains_no_key_in_treap(self):
node = TreapNode("cdd", 30)
node = node.get_with_inserted("ab", 23)
return self.assertFalse("zero" in node)
def test_contains_one_node(self):
node = TreapNode(10, "10")
return self.assertTrue(10 in node)
def get_big_treap_node():
node = TreapNode(-1, "-1")
keys = set([random.randint(0, 500) for _ in range(100)])
for key in keys:
node = node.get_with_inserted(key, f"{key}")
return node