Endre Moen
Endre Moen

Reputation: 784

delete arbitrary item from heap in python

Is there a binary heap implementation where I can pop other elements than the root in log n time?

I use heapq - but heap.index( wKeys ) in

heap.pop( heap.index( wKeys ) )

is very slow. I need a binary heap for my problem - where I sometime use

heapq.heappop(heap)

but also need to pop other elements than at the top from the heap. So a binary heap like heapq imlementation should do that but I havnt found a binary search method. I also looked at treap (http://stromberg.dnsalias.org/~strombrg/treap/) and cannot find a method like this here either.

Upvotes: 2

Views: 3948

Answers (2)

Endre Moen
Endre Moen

Reputation: 784

Ive modified the implementation of heapq by adding a parameter to heappop(), and heappush() - which is heapIndex. That takes a dictionary of {item: index} and updates the heapIndex as it pops or pushes into heap.

Ive also added a new method heappop_arbitrary() which deletes an arbitrary element and updates the heapIndex

The code is available here: https://github.com/emoen/heapq_with_index

Ive renamed the methods heappop(),heappush() to heappop2(), heappush2() to avoid confusion with the original methods.

I havent implemented this for any of the other functions available in heapq.

Edit: please star if you can use the repo :)

Upvotes: 1

Isan Sahoo
Isan Sahoo

Reputation: 404

class RemoveHeap:
    def __init__(self):
        self.h = []
        self.track = collections.defaultdict(collections.deque)
        self.counter = itertools.count()

    def insert_item(self, val):
        count = next(self.counter)
        item = [val, count, 'active']
        self.track[val].append(item)
        heapq.heappush(self.h, item)

    def delete_item(self, val):
        if val in self.track:
            items = self.track[val]
            for item in items:
                if item[2] == 'active':
                    item[2] = 'deleted'
                    break

    def pop_item(self):
        while len(self.h) > 0:
            item = heapq.heappop(self.h)
            item_track = self.track[item[0]]
            item_track.popleft()
            if len(item_track) == 0:
                del self.track[item[0]]
            else:
                self.track[item[0]] = item_track
            if item[2] == 'active':
                return item[0]

    def peek_item(self):
        item = self.h[0]
        if item[2] == 'deleted':
            x = self.pop_item()
            self.insert_item(x)
            return x
        return item[0]

Upvotes: 0

Related Questions