Added __len__, __contains__, clear methods.

Fixed issue with not unique heap keys

Author: yusinv

CHANGELOG.rst

@@ -14,7 +14,10 @@ Changelog
 * Initial implementation.
 
 
-0.0.2 (2024-03-24)
+0.0.2 (2024-03-26)
 ------------------
 
 * Added sorted_query method.
+* Added __len__ and __contains__ methods.
+* Added clear method.
+* Fixed issue with not unique heap keys

docs/reference/priority_search_tree.rst

@@ -7,3 +7,4 @@ priority_search_tree
 
 .. automodule:: priority_search_tree
     :members:
+    :special-members:

src/priority_search_tree/__init__.py

@@ -40,20 +40,47 @@ class PrioritySearchTree:
     """
 
     def _push_down(self, node: Node, value: _V) -> None:
-        while node != Node.NULL_NODE:
-            if node.heap_value is None:
-                node.heap_value = value
-                return
+        if node == Node.NULL_NODE:
+            return
+        if node.heap_value is None:
+            node.heap_value = value
+            return
 
-            if self.heap_key(value) > self.heap_key(node.heap_value):
-                node.heap_value, value = value, node.heap_value
-            elif self.heap_key(value) == self.heap_key(node.heap_value):
-                node.placeholder = False
+        if self.tree_key(node.heap_value) < self.tree_key(node.tree_value):
+            self._push_down(node.left, node.heap_value)
+        else:
+            self._push_down(node.right, node.heap_value)
 
-            if self.tree_key(value) < self.tree_key(node.tree_value):
-                node = node.left
-            else:
-                node = node.right
+        if node.placeholder:
+            node.placeholder = False
+        else:
+            node.heap_value = value
+
+    def _sieve_down(self, node: Node, value: _V) -> None:
+
+        if node.heap_value is None:
+            node.heap_value = value
+            return
+
+        if node.placeholder:
+            node.placeholder = False
+            return
+
+        heap_key_value = self.heap_key(value)
+        tree_key_value = self.tree_key(value)
+
+        if heap_key_value > self.heap_key(node.heap_value):
+            self._push_down(node, value)
+            return
+
+        if heap_key_value == self.heap_key(node.heap_value) and tree_key_value < self.tree_key(node.heap_value):
+            self._push_down(node, value)
+            return
+
+        if tree_key_value < self.tree_key(node.tree_value):
+            self._sieve_down(node.left, value)
+        else:
+            self._sieve_down(node.right, value)
 
     def _push_up(self, node: Node) -> None:
         vl, vr = None, None
@@ -65,7 +92,7 @@ def _push_up(self, node: Node) -> None:
             vr = node.right.heap_value
 
         if vl and vr:
-            if self.heap_key(vl) > self.heap_key(vr):
+            if self.heap_key(vl) >= self.heap_key(vr):
                 node.heap_value = vl
                 self._push_up(node.left)
             else:
@@ -230,7 +257,7 @@ def add(self, value: _V) -> None:
             new_internal_node.heap_value = prev.heap_value
             prev.placeholder = True
 
-        self._push_down(self._root, value)
+        self._sieve_down(self._root, value)
         self._fix_insert(new_leaf_node)
         self._len += 1
 
@@ -246,17 +273,19 @@ def remove(self, value: _V) -> None:
 
         Complexity:
             O(log(N)) where **N** is number of items in PST
+
+        Note:
+            this function is using ``tree_key(value)`` to compare the items
         """
         node = self._root
         value_tree_key = self.tree_key(value)
-        value_heap_key = self.heap_key(value)
         heap_node = None
         tree_node = None
         leaf_node = None
 
         while node != Node.NULL_NODE:
             leaf_node = node
-            if heap_node is None and self.tree_key(node.heap_value) == value_tree_key and self.heap_key(node.heap_value) == value_heap_key:
+            if heap_node is None and self.tree_key(node.heap_value) == value_tree_key:
                 heap_node = node
             if tree_node is None and self.tree_key(node.tree_value) == value_tree_key:
                 tree_node = node
@@ -286,7 +315,7 @@ def remove(self, value: _V) -> None:
             cut_node = leaf_node.parent
             fix_node = leaf_node.parent.right
 
-        self._push_down(cut_node, cut_node.heap_value)
+        self._push_down(cut_node, None)
         self._transplant(cut_node, fix_node)
 
         if cut_node.color == 0:
@@ -535,5 +564,53 @@ def _rotate_left(self, x: Node) -> None:
         y.set_left(x)
         self._push_up(x)
 
-    def __len__(self):
+    def __len__(self) -> int:
+        """
+        Implements the built-in function len()
+
+        Returns:
+            int: Number of items in PST.
+
+        Complexity:
+            O(1)
+        """
         return self._len
+
+    def __contains__(self, value) -> bool:
+        """
+        Implements membership test operator.
+
+        Args:
+            value: Value to test for membership
+
+        Returns:
+            bool:  ``True`` if value is in ``self``, ``False`` otherwise.
+
+        Complexity:
+            O(log(N)) where **N** is number of items in PST
+
+        Note:
+            this function is using ``tree_key(value)`` to compare the items
+        """
+        value_tree_key = self.tree_key(value)
+
+        node = self._root
+        while node != Node.NULL_NODE:
+            if value_tree_key < self.tree_key(node.tree_value):
+                node = node.left
+            elif value_tree_key == self.tree_key(node.tree_value):
+                return True
+            else:
+                node = node.right
+
+        return False
+
+    def clear(self) -> None:
+        """
+        Removes **all** items from PST.
+
+        Complexity:
+            O(1)
+        """
+        self._root = Node.NULL_NODE
+        self._len = 0

tests/test_priority_search_tree.py

@@ -12,6 +12,7 @@
 def test_empty_pst():
     pst = PrioritySearchTree()
     assert not pst
+    assert (1, 1) not in pst
     result = pst.query((0, 0), (1, 2), (1, 1))
     assert len(result) == 0
     with pytest.raises(ValueError, match="Value not found:"):
@@ -67,6 +68,20 @@ def test_heap_get_max():
     assert pst.heap_get_max() == (5, 5)
 
 
+def test_contains():
+    items = [(1, 1), (2, 2), (3, 4), (4, 4)]
+    pst = PrioritySearchTree(items, tree_key=lambda x: x[0])
+    for itm in items:
+        assert itm in pst
+
+    # key_tree and heap_tree not in pst
+    assert (5, 5) not in pst
+    # key_tree not in pst and heap_key is
+    assert (5, 4) not in pst
+    # key_tree in pst heap_key is not
+    assert (1, 5) in pst
+
+
 def test_large_pst():
     pst = PrioritySearchTree(LARGE_PST_INITIAL_DATA)
     assert len(pst) == len(LARGE_PST_INITIAL_DATA)
@@ -134,6 +149,35 @@ def test_unique_tree_key():
         PrioritySearchTree([(1, 1), (1, 1)])
 
 
+def test_not_unique_heap_keys():
+    pst = PrioritySearchTree()
+    for i in range(100, -1, -1):
+        pst.add((i, 5))
+    assert_rb_tree(pst._root)
+    result = pst.sorted_query((0, 0), (1000, 0), (0, 0))
+    assert len(result) == len(pst)
+    for i, r in enumerate(result):
+        assert r[0] == i
+
+    pst.clear()
+    for i in range(100):
+        pst.add((i, 5))
+    assert_rb_tree(pst._root)
+    result = pst.sorted_query((0, 0), (1000, 0), (0, 0))
+    assert len(result) == len(pst)
+    for i, r in enumerate(result):
+        assert r[0] == i
+
+    pst.clear()
+    for itm in [(0, 5), (1, 5), (2, 5), (3, 5), (4, 5), (5, 5), (6, 5), (7, 5), (8, 5)]:
+        pst.add(itm)
+    assert_rb_tree(pst._root)
+    result = pst.sorted_query((0, 0), (1000, 0), (0, 0))
+    assert len(result) == len(pst)
+    for i, r in enumerate(result):
+        assert r[0] == i
+
+
 def test_custom_keys():
     class Point:
         def __init__(self, x: int, y: int):
@@ -155,8 +199,7 @@ def __init__(self, x: int, y: int):
     pst.remove(Point(2, 2))
 
     # same tree_key different hash_key
-    with pytest.raises(ValueError, match="Value not found:"):
-        pst.remove(Point(4, 1))
+    pst.remove(Point(4, 1))
 
     # different tree_key same hash_key
     with pytest.raises(ValueError, match="Value not found:"):
@@ -170,5 +213,4 @@ def __init__(self, x: int, y: int):
 
     assert pst.heap_pop().y == 6
     assert pst.heap_pop().y == 6
-    assert pst.heap_pop().y == 4
     assert pst.heap_pop().y == 1