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:
- implement your own heap with a removal operation,
- augment a standard heap with lazy removals.
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