CODEWITHSUNDEEP
pythonperformanceruntime
Trim
Trim

Reputation: 373

Python heapq module, heapify method on an object

Since I'm trying to be efficient in this program I'm making, I thought I'd use the built in heapq module in python, but some of my objects have multiple attributes, like name and number. Is there a way to use the heapify method to heapify my objects based on a certain attribute? I don't see anything in the documentation.

Upvotes: 1

Views: 1142

Answers (2)

rocksteady
rocksteady

Reputation: 2560

@vsekhar and @ all the others wondering about the accepted answer.

Assumption:

class SomeObject():
    def __init__(self,name, number):
        self.name = name
        self.number = number

a_list = []
obj_1 = SomeObject("tim", 12)
obj_2 = SomeObject("tom", 13)

Now, instead of creating a heap with the objects only as elements:

heapq.heappush(a_list, obj_1)
heapq.heappush(a_list, obj_2)

you actually want to create the heap with a tuple of 2 values as heap elements - The idea is to have the attribute you want to sort with as first value of the tuple and the object (as before) as the second element of the tuple:

# Sort by 'number'.
heapq.heappush(a_list, (obj_1.number, obj_1))
heapq.heappush(a_list, (obj_2.number, obj_2))

The heap considers this first value of the tuple as the value to sort by.

  • In case the element pushed to the heap is not of a simple data type like int or str, the underlying implementation needs to know how to compare elements.
  • If the element is an iterable the first element is considered to contain the sort value.
  • Have a look at the examples here: https://docs.python.org/3/library/heapq.html#basic-examples (search for tuple)

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

Another option might be to make comparison work with your custom class - This can be implemented so the object itself can be used as the heap element (as in the first example).

class SomeObject():
    def __init__(self,name, number):
        self.name = name
        self.number = number
    def __eq__(self, obj): 
        return self.number == obj.number 
    def __lt__(self, obj): 
        return self.number < obj.number 
    def __hash__(self): 
        return hash(self.number)

This way you can create the heap with the objects only as elements:

heapq.heappush(a_list, obj_1)
heapq.heappush(a_list, obj_2)

Upvotes: 1

Trim
Trim

Reputation: 373

Right after I posted, I figured you could make a list of the objects by the attribute needed before using heapify which would take O(n) linear time. This wouldn't affect the runtime of heapify or other heapq methods.

Upvotes: 1

Related Questions