Python 3优先级队列：如何以O(log(n))复杂度查找并删除指定值元素

Great question! The built-in queue.PriorityQueue is built on top of Python's heapq module, which doesn't natively support fast lookup or deletion of arbitrary elements—heaps are optimized for push/pop operations at the top, not random access. But we can add a lazy deletion mechanism alongside a secondary index to achieve amortized O(log n) time for these operations.

How It Works #

The core idea is:

• Don't immediately remove elements from the heap (which would take O(n) time to locate first).
• Instead, mark elements as "deleted" in a separate dictionary, and clean up invalid elements from the heap only when we try to access the top element.
• Use a unique token to handle duplicate values (so we can distinguish between multiple elements with the same priority).

Implementation Code #

Here's an enhanced priority queue that adds the functionality you need:

from queue import PriorityQueue
from collections import defaultdict

class PqElement(object):
    def __init__(self, value: int, token: int):
        self.val = value
        self.token = token  # Unique identifier to handle duplicate values
    # Maintain your max-heap behavior
    def __lt__(self, other):
        return self.val > other.val
    def __repr__(self):
        return f'PQE:{self.val}(token:{self.token})'

class EnhancedPriorityQueue:
    def __init__(self):
        self._pq = PriorityQueue()
        self._entry_index = defaultdict(list)  # Maps value to list of active tokens
        self._token_counter = 0  # Generates unique tokens for each element
    
    def put(self, value: int):
        """Add an element to the priority queue (O(log n))"""
        token = self._token_counter
        self._token_counter += 1
        elem = PqElement(value, token)
        self._pq.put(elem)
        self._entry_index[value].append(token)
    
    def _cleanup_invalid_elements(self):
        """Remove any marked-deleted elements from the heap top (amortized O(log n))"""
        while not self._pq.empty():
            top_elem = self._pq.queue[0]
            # Check if this element's token is still in the active index
            if top_elem.token not in self._entry_index[top_elem.val]:
                self._pq.get()  # Remove invalid element from heap
            else:
                break  # Heap top is valid, stop cleanup
    
    def get(self):
        """Pop the highest-priority element (amortized O(log n))"""
        self._cleanup_invalid_elements()
        if self._pq.empty():
            raise IndexError("Cannot get from empty priority queue")
        
        elem = self._pq.get()
        # Remove the token from our active index
        self._entry_index[elem.val].remove(elem.token)
        # Clean up the index key if no elements are left for this value
        if not self._entry_index[elem.val]:
            del self._entry_index[elem.val]
        return elem
    
    def remove(self, value: int) -> bool:
        """Delete one occurrence of the specified value (amortized O(log n))"""
        # Check if the value exists in our active elements
        if value not in self._entry_index or not self._entry_index[value]:
            return False
        
        # Mark the first occurrence as deleted by removing its token from the index
        token_to_remove = self._entry_index[value].pop(0)
        if not self._entry_index[value]:
            del self._entry_index[value]
        
        # Optional: Clean up the heap immediately (or wait for next get/peek)
        self._cleanup_invalid_elements()
        return True
    
    def peek(self):
        """Get the highest-priority value without removing it (O(1) amortized)"""
        self._cleanup_invalid_elements()
        return self._pq.queue[0].val if not self._pq.empty() else None
    
    def qsize(self):
        """Get the count of active elements (O(1))"""
        return sum(len(tokens) for tokens in self._entry_index.values())
    
    def empty(self):
        """Check if the queue has any active elements (O(1))"""
        return self.qsize() == 0

Key Details #

• Lazy Deletion: When you call remove(), we just remove the element's token from our index—we don't touch the heap yet. Invalid elements are only cleaned up when we try to access the heap top (via get() or peek()), which ensures each element is processed at most once, leading to amortized O(log n) time.
• Duplicate Handling: The unique token ensures we can track and delete individual elements even if they have the same priority value.
• Time Complexity:
  • put(): O(log n) (same as original, plus O(1) dictionary operation)
  • remove(): Amortized O(log n) (O(1) index update plus deferred heap cleanup)
  • get(): Amortized O(log n) (heap pop plus deferred cleanup)

Usage Notes #

• Avoid directly accessing _pq.queue anymore—use peek() to get the top value, as the raw heap may contain invalid, marked-deleted elements.
• The qsize() method returns the count of active elements, not the size of the underlying heap (which may include stale elements).

内容的提问来源于stack exchange，提问作者S M Abrar Jahin