gh-136898: feat(collections)-Add HeapDict class

Lib/collections/__init__.py

@@ -1606,3 +1606,139 @@ def upper(self):
 
     def zfill(self, width):
         return self.__class__(self.data.zfill(width))
+
+
+################################################################################
+### HeapDict
+################################################################################
+
+import collections.abc
+import heapq
+import sys
+
+# Use string forward references for type annotations to avoid circular imports
+if sys.version_info >= (3, 7):
+    from __future__ import annotations
+
+class HeapDict(collections.abc.MutableMapping):
+    """Dictionary that maintains heap property based on values.
+
+    HeapDict combines the functionality of a dictionary with a heap,
+    providing efficient access to key-value pairs while maintaining
+    a heap property for priority-based operations.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Initialize a new HeapDict with optional initial values."""
+        self._dict = {}  # Maps keys to values
+        self._heap = []  # List of (value, key, counter)
+        self._counter = 0  # Used to break ties for values that compare equal
+        self._removed_keys = set()  # Track removed keys for lazy deletion
+
+        # Add initial items
+        if args or kwargs:
+            self.update(*args, **kwargs)
+
+    def __setitem__(self, key, value):
+        """Set a key-value pair, maintaining heap property."""
+        if key in self._dict:
+            self.update_priority(key, value)
+        else:
+            self._dict[key] = value
+            count = self._counter
+            self._counter += 1
+            heapq.heappush(self._heap, (value, key, count))
+
+    def __getitem__(self, key):
+        """Get value by key."""
+        return self._dict[key]
+
+    def __delitem__(self, key):
+        """Remove a key-value pair."""
+        if key not in self._dict:
+            raise KeyError(key)
+
+        # Mark the key as removed
+        self._removed_keys.add(key)
+        del self._dict[key]
+
+    def __iter__(self):
+        """Iterate through keys."""
+        return iter(self._dict)
+
+    def __len__(self):
+        """Return the number of items."""
+        return len(self._dict)
+
+    def __repr__(self):
+        """Return string representation."""
+        return f"{self.__class__.__name__}({dict(self.items())})"
+
+    def _clean_heap(self):
+        """Clean the heap by removing marked items."""
+        if len(self._removed_keys) > len(self._heap) // 2:
+            # If too many removed items, rebuild the heap
+            new_heap = [(v, k, c) for v, k, c in self._heap if k in self._dict]
+            heapq.heapify(new_heap)
+            self._heap = new_heap
+            self._removed_keys.clear()
+
+    def popmin(self):
+        """Remove and return the (key, value) pair with minimum value."""
+        if not self._dict:
+            raise KeyError("popmin from an empty HeapDict")
+
+        # Skip items that were already removed
+        while self._heap:
+            value, key, _ = heapq.heappop(self._heap)
+            if key not in self._removed_keys and key in self._dict:
+                del self._dict[key]
+                return key, value
+
+        raise RuntimeError("Heap is inconsistent with dictionary")
+
+    def peekmin(self):
+        """Return the (key, value) pair with minimum value without removing it."""
+        if not self._dict:
+            raise KeyError("peekmin from an empty HeapDict")
+
+        # Skip items that were already removed
+        while self._heap:
+            value, key, _ = self._heap[0]
+            if key not in self._removed_keys and key in self._dict:
+                return key, value
+            heapq.heappop(self._heap)
+
+        raise RuntimeError("Heap is inconsistent with dictionary")
+
+    def update_priority(self, key, new_value):
+        """Update the value/priority of an existing key."""
+        if key not in self._dict:
+            raise KeyError(key)
+
+        # Update the dictionary
+        self._dict[key] = new_value
+        self._removed_keys.add(key)
+        count = self._counter
+        self._counter += 1
+        heapq.heappush(self._heap, (new_value, key, count))
+
+        if len(self._removed_keys) > len(self._heap) // 2:
+            self._clean_heap()
+
+# Add type hints for older Python versions or static type checkers
+if sys.version_info < (3, 7) or typing.TYPE_CHECKING:
+    from typing import Any, Dict, Iterator, List, Optional, Tuple, TypeVar, Union
+    K = TypeVar('K')
+    V = TypeVar('V')
+    HeapDict.__annotations__ = {
+        '_dict': Dict[K, V],
+        '_heap': List[Tuple[V, K, int]],
+        '__setitem__': None,
+        '__getitem__': None,
+        '__delitem__': None,
+        '__iter__': Iterator[K],
+        'popmin': Tuple[K, V],
+        'peekmin': Tuple[K, V],
+        'update_priority': None
+    }

Misc/NEWS.d/next/Library/2025-07-21-15-29-04.gh-issue-136898.YAz31U.rst

@@ -0,0 +1 @@
+Add HeapDict in collection