Heapq with custom predicate in Python

Prerequisite: heapq module

The heapq module has several functions that take the list as a parameter and arranges it in a min-heap order. The problem with these functions is they expect either a list or a list of tuples as a parameter. They do not support comparisons between any other iterable or objects. For example, consider a dictionary that has to be maintained in heap.

import heapq as hq

my_dict={'a':'apple', 'b':'ball', 'c': 'cat'}
hq.heapify(my_dict)

print(my_dict)

Output 

TypeError: heap argument must be a list

The heapify() function expects the parameter to be a list. So if we consider a list of dictionaries, look below what happens.

import heapq as hq

my_dict=[{'a':'apple'}, {'b':'ball'}, {'c': 'cat'}]

hq.heapify(my_dict)

print(my_dict)

Output:

TypeError: '<' not supported between instances of 'dict' and 'dict' 
 

Thus, we cannot compare two dictionaries using the heapq module. Sometimes we may have to compare objects of a class and maintain them in a heap. The comparison between such objects is also not feasible with this module.

#import module
import heapq as hq

# the dictionary to be as heap
my_dict = {'z': 'zebra', 'b': 'ball', 'w': 'whale', 
           'a': 'apple', 'm': 'monkey', 'c': 'cat'}

# conversion to tuple
my_list = [(k, v) for k, v in my_dict.items()]

print("Before organizing as heap :", my_list)

# arrange as min-heap
hq.heapify(my_list)

print("After organizing as heap :", my_list)

# re convert to dictionary
my_dict = dict(my_list)

print("Resultant dictionary :", my_dict)

This article discusses how to overcome the above-said issues.

Customizing the sort in heapq

The heapq module functions can take either a list of items or a list of tuples as a parameter. Thus, there are two ways to customize the sorting process:

1.  Convert the iterable to a list of tuples/list for comparison.
2.  Write a wrapper class that overrides '<' operator.

Conversion to list of items

This method is simple and can be used for solving dictionary comparison problems. The dictionary items can be converted into a list of tuples and then passed to the heapify method.

Example 1: a simple dictionary

#import module
import heapq as hq

# the dictionary to be as heap
my_dict = {'z': 'zebra', 'b': 'ball', 'w': 'whale',
           'a': 'apple', 'm': 'monkey', 'c': 'cat'}

# conversion to tuple
my_list = [(k, v) for k, v in my_dict.items()]

print("Before organizing as heap :", my_list)

# arrange as min-heap
hq.heapify(my_list)

print("After organizing as heap :", my_list)

# re convert to dictionary
my_dict = dict(my_list)

print("Resultant dictionary :", my_dict)

Output:

Before organizing as heap : [('z', 'zebra'), ('b', 'ball'), ('w', 'whale'), ('a', 'apple'), ('m', 'monkey'), ('c', 'cat')] 
After organizing as heap : [('a', 'apple'), ('b', 'ball'), ('c', 'cat'), ('z', 'zebra'), ('m', 'monkey'), ('w', 'whale')] 
Resultant dictionary : {'a': 'apple', 'b': 'ball', 'c': 'cat', 'z': 'zebra', 'm': 'monkey', 'w': 'whale'}

Example 2: a list of dictionaries

#import module
import heapq as hq

# the list of dictionaries to be as heap
my_dict = [{'z': 'zebra'}, {'b': 'ball'}, {'w': 'whale'},
           {'a': 'apple'}, {'m': 'monkey'}, {'c': 'cat'}]

# conversion to tuple
my_list = [(k, v) for i in my_dict for k, v in i.items()]

print("Before organizing as heap :", my_list)

# arrange as min-heap
hq.heapify(my_list)

print("After organizing as heap :", my_list)

# re convert to dictionary
my_dict = dict(my_list)

print("Resultant dictionary :", my_dict)

Output:

Before organizing as heap : [('z', 'zebra'), ('b', 'ball'), ('w', 'whale'), ('a', 'apple'), ('m', 'monkey'), ('c', 'cat')] 
After organizing as heap : [('a', 'apple'), ('b', 'ball'), ('c', 'cat'), ('z', 'zebra'), ('m', 'monkey'), ('w', 'whale')] 
Resultant dictionary : {'a': 'apple', 'b': 'ball', 'c': 'cat', 'z': 'zebra', 'm': 'monkey', 'w': 'whale'}

The above methods can be used for a dictionary with any data type. 

Using a wrapper class

Consider a situation where the objects of a class have to be maintained in a min-heap.  For example, let us consider a class that has attributes like 'name', 'designation', 'yos'(years of service), 'salary'. The objects of this class have to be maintained in min-heap based on 'yos' (years of service). 

Here, we override the relational operator '<' such that it compares the years of service of each employee and returns true or false. Based on the returned boolean value, heapq module arranges the objects in min-heap order.

# import required module
import heapq as hq

# class definition
class employee:

  # constructor
    def __init__(self, n, d, yos, s):
        self.name = n
        self.des = d
        self.yos = yos
        self.sal = s

  # function for customized printing
    def print_me(self):
        print("Name :", self.name)
        print("Designation :", self.des)
        print("Years of service :", str(self.yos))
        print("salary :", str(self.sal))

   # override the comparison operator
    def __lt__(self, nxt):
        return self.yos < nxt.yos


# creating objects
e1 = employee('Anish', 'manager', 3, 24000)
e2 = employee('kathy', 'programmer', 2, 15000)
e3 = employee('Rina', 'Analyst', 5, 30000)
e4 = employee('Vinay', 'programmer', 1, 10000)

# list of employee objects
emp = [e1, e2, e3, e4]

# converting to min-heap
# based on yos
hq.heapify(emp)

# printing the results
for i in range(0, len(emp)):
    emp[i].print_me()
    print()

Name : Vinay
Designation : programmer
Years of service : 1
salary : 10000

Name : kathy
Designation : programmer
Years of service : 2
salary : 15000

Name : Rina
Designation : Analyst
Years of service : 5
salary : 30000

Name : Anish
Designation : manager
Years of service : 3
salary : 24000