utils.py

import itertools
from typing import Any
from heapq import heappush, heappop


class PriorityQueue:
    REMOVED = '<removed-task>'  # placeholder for a removed task

    def __init__(self):
        self.pq = []  # list of entries arranged in a heap
        self.entry_finder = {}  # mapping of tasks to entries
        self.counter = itertools.count()  # unique sequence count as tie break

    def insert(self, item: Any, priority: Any = 0, count: Any = None):
        """Add a new task or update the priority of an existing task"""
        if item in self.entry_finder:
            self.remove_item(item)

        # Additional manual set tie break
        if count is not None:
            entry = [priority, count, next(self.counter), item]
        else:
            entry = [priority, next(self.counter), item]

        self.entry_finder[item] = entry
        heappush(self.pq, entry)

    def remove_item(self, item: Any):
        """Mark an existing task as REMOVED.  Raise KeyError if not found."""
        entry = self.entry_finder.pop(item)
        entry[-1] = self.REMOVED

    def delete_min(self) -> Any:
        """Remove and return the lowest priority task. Raise KeyError if empty."""
        while self.pq:
            item = heappop(self.pq)[-1]
            if item is not self.REMOVED:
                del self.entry_finder[item]
                return item
        raise KeyError('pop from an empty priority queue')

    def peak_min(self) -> Any:
        """Return minimum item without removing it from the queue."""
        while self.pq:
            item = self.pq[0][-1]
            if item is not self.REMOVED:
                return item
            else:
                heappop(self.pq)
        raise KeyError('Empty queue.')