TryAlgo

Heap allowing to remove items

Maintain a set, allowing to add or remove elements and to query the smallest element.

Heap

Such a data structure is usually implemented by a heap. Heaps are binary trees, where every node stores a smaller value than its descendants. The heaps available in the standard libraries usually only allow to insert items.

Hence you have two possibilities:

Lazy removals

The later option is much simpler to implement. Simply mark elements to be removed, and remove them only once they appear as the top element of the heap. Since the heap can store multiple copies of the same item, we use a counter to mark items. It simply stores how many copies of a given item have to removed.

The price for this simple implementation is a larger worst case complexity of the operations. If your heap contains n elements plus k elements still to be removed, then the operations have time complexity O(log(n + k)). Implementing your own heap with a real removal operator, allows to achieve complexity O(log n) per operation. However for most problems in programming competitions, the lazy removals are good enough.

Implementation in Python

class lazyheap:
    """Maintains a heap where elements can be removed.
    Elements do not have to be distinct.
    Removals are done in lazy manner, namely only when seen at top.
    toremove[v] = how many times v has yet to be removed.
    """
    def __init__(self):
        self.h = []  # the actual heap
        self.n = 0   # number of (non removed) items in heap
        self.sum = 0 # their sum
        self.toremove = Counter()
    
    def push(self, item):
        heappush(self.h, item)
        self.n += 1
        self.sum += item

    def remove(self, item):
        # just mark for later actual removal
        self.toremove[item] += 1    
        self.n -= 1
        self.sum -= item

    def top(self):
        """returns smallest (non removed) item of heap"""
        x = self.h[0]
        while self.toremove[x] > 0:
            self.toremove[x] -= 1
            heappop(self.h)
            x = self.h[0]
        return x

    def pop(self):
        """ removes and returns smallest element """
        x = self.top()
        heappop(self.h)
        self.n -= 1
        self.sum -= x
        return x