Question

Python has heapq module which implements heap data structure and it supports some basic operations (push, pop).

How to remove i-th element from the heap in O(log n)? Is it even possible with heapq or do I have to use another module?

Note, there is an example at the bottom of the documentation:
http://docs.python.org/library/heapq.html
which suggest a possible approach – this is not what I want. I want the element to remove, not to merely mark as removed.

Answer 1

You can remove the i-th element from a heap quite easily:

h[i] = h[-1]
h.pop()
heapq.heapify(h)

Just replace the element you want to remove with the last element and remove the last element then re-heapify the heap. This is O(n), if you want you can do the same thing in O(log(n)) but you’ll need to call a couple of the internal heapify functions, or better as larsmans pointed out just copy the source of _siftup/_siftdown out of heapq.py into your own code:

h[i] = h[-1]
h.pop()
if i < len(h):
    heapq._siftup(h, i)
    heapq._siftdown(h, 0, i)

Note that in each case you can’t just do h[i] = h.pop() as that would fail if i references the last element. If you special case removing the last element then you could combine the overwrite and pop.

Note that depending on the typical size of your heap you might find that just calling heapify while theoretically less efficient could be faster than re-using _siftup/_siftdown: a little bit of introspection will reveal that heapify is probably implemented in C but the C implementation of the internal functions aren’t exposed. If performance matter to you then consider doing some timing tests on typical data to see which is best. Unless you have really massive heaps big-O may not be the most important factor.

Edit: someone tried to edit this answer to remove the call to _siftdown with a comment that:

_siftdown is not needed. New h[i] is guaranteed to be the smallest of the old h[i]’s children, which is still larger than old h[i]’s parent
(new h[i]’s parent). _siftdown will be a no-op. I have to edit since I
don’t have enough rep to add a comment yet.

What they’ve missed in this comment is that h[-1] might not be a child of h[i] at all. The new value inserted at h[i] could come from a completely different branch of the heap so it might need to be sifted in either direction.

Also to the comment asking why not just use sort() to restore the heap: calling _siftup and _siftdown are both O(log n) operations, calling heapify is O(n). Calling sort() is an O(n log n) operation. It is quite possible that calling sort will be fast enough but for large heaps it is an unnecessary overhead.

Edited to avoid the issue pointed out by @Seth Bruder. When i references the end element the _siftup() call would fail, but in that case popping an element off the end of the heap doesn’t break the heap invariant.