CODEWITHSUNDEEP
pythonheap
Oblomov
Oblomov

Reputation: 9635

Define heap key for an array of tuples

A simple example for the usage of the python heap implementation is

from heapq import heappush, heappop
heap = []
data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
for item in data:
    heappush(heap, item)

In a more complicated scenario, I have an array of tuples like

tuples = [(5,"foo",True),(2,"bar", False),(8,"foobar",True)]

and want to use the first entry of each tuple as heap key, i.e. the tuples should be sorted according to the number in the tuples.

How can I do that?

Upvotes: 13

Views: 11209

Answers (1)

Chris
Chris

Reputation: 22953

You can simply use the tuple as they are. The Python documentation explicitly makes note of such as usage:

Heap elements can be tuples. This is useful for assigning comparison values (such as task priorities) alongside the main record being tracked:

>>> h = []
>>> heappush(h, (5, 'write code'))
>>> heappush(h, (7, 'release product'))
>>> heappush(h, (1, 'write spec'))
>>> heappush(h, (3, 'create tests'))
>>> heappop(h)
(1, 'write spec')

Simply push the tuples to the heap, and pop them off when needed:

>>> from heapq import heappush, heappop
>>> 
>>> heap = []
>>> tuples = [(5,"foo",True),(2,"bar", False),(8,"foobar",True)] 
>>> 
>>> for tup in tuples:
...     heappush(heap, tup)
... 
>>> heappop(heap)
(2, 'bar', False)

Because the implementation for heap uses default sorting for tuples

while pos > startpos:
    ...
    if newitem < parent:
        ...
    ...
...

and Python sorts tuples element-wise, ensure the objects by which you want the tuples to be sorted come first.

Upvotes: 16

Related Questions