elllot
elllot

Reputation: 373

Python overriding heapq comparison when using tuples

So I understand that you can use tuples instead of straight values in using the heapq module for creating a heap that essentially contains (key,value) pairs.

I'm also aware that you can override the __lt__ comparison operator for the heapq module to have your own comparison in creating and maintaining the heap.

Is there a way to override this operator so that it compares both values of the tuple? For instance, I'm trying to have heapq compare the second values of the tuple if the first are equal with something along the lines of:

def __lt__(self, other):
    return self[0] < other[0] if self[0] != other[0] else other[1] < self[1]

The latter half other[1] < self[1] is reversed on purpose as the logic is that the tuple with a greater second value should be considered the lesser of the two.

Thanks for any insight in advance!

Edit: I think I was confused by the difference between the terminologies "overload" and "override". Override was what I was referring to.

Upvotes: 1

Views: 2594

Answers (1)

Lawrence D&#39;Oliveiro
Lawrence D&#39;Oliveiro

Reputation: 2774

Yes, that will work. To clarify the overall process for using a custom class with heapq, here are some extracts from a test script I wrote to try to get to grips with asyncio. To keep track of timers that go off at different times, I defined this class:

class TimerTask :
    "actions to be invoked at specific times."

    def __init__(self, when, action) :
        self.when = when
        self.action = action
    #end __init__

    def __lt__(a, b) :
        return \
            a.when < b.when
    #end __lt__

#end TimerTask

A line like this puts a TimerTask on the pending queue:

heapq.heappush(self.pending, self.TimerTask(when, trigger_sleep_done))

Then this sequence waits for the earliest pending timer to come due and invokes its corresponding action:

try :
    until = self.pending[0].when
except IndexError :
    # nothing more for event loop to do
    break
#end try
now = time.time()
if until > now :
    time.sleep(until - now)
#end if
heapq.heappop(self.pending).action()

Upvotes: 1

Related Questions