How to Use Python heapq Module

09/13/2021

Contents

In this article, you will learn how to use Python heapq module.

Python heapq Module

Python’s heapq module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm. A heap is a specialized tree-based data structure where the parent nodes are always greater than or equal to their child nodes, and the minimum element is always at the root. This makes it useful for implementing priority queues where the smallest element needs to be accessed quickly.

Here are some steps to use the heapq module in Python:

  1. Import the heapq module by adding “import heapq” at the beginning of your Python script.

    import heapq
  2. Create a list of elements that you want to add to the heap. The heapq module works on any iterable, so you can use a list, tuple, or any other iterable.
  3. Use the heapq.heapify() function to convert the list into a heap. This function rearranges the elements in the list so that they satisfy the heap property.

    a = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
    heapq.heapify(a)
    
  4. To add an element to the heap, use the heapq.heappush() function. This function adds the element to the heap and maintains the heap property.

    heapq.heappush(a, 7)
  5. To remove the smallest element from the heap, use the heapq.heappop() function. This function removes the root element (i.e., the smallest element) from the heap and returns it.

    smallest = heapq.heappop(a)
  6. To access the smallest element in the heap without removing it, use the heapq.nsmallest() function. This function returns the smallest n elements from the heap without modifying the heap itself.

    smallest = heapq.nsmallest(1, a)

By using these functions, you can easily implement a priority queue in Python using the heapq module.

Here are some additional details and examples to help you better understand the Python heapq module:

  • The heapq module implements the heap data structure using a regular Python list. The elements in the list are ordered in such a way that the minimum element is always at the root of the heap. This means that the first element in the list (i.e., index 0) is always the smallest.
  • The heapq module provides several other functions for working with heaps, such as heapq.heappushpop() (which pushes a new element onto the heap and then pops the smallest element), heapq.heapreplace() (which pops the smallest element and then pushes a new element onto the heap), heapq.nlargest() (which returns the largest n elements in the heap), and heapq.merge() (which merges multiple sorted iterables into a single sorted iterable using a heap).
  • The heapq module can be used with custom objects by providing a key function that returns the value to use for comparison. For example, suppose you have a list of tuples where each tuple represents a person’s name and age. You can use heapq.nsmallest() to find the youngest person in the list as follows:

    people = [('Alice', 25), ('Bob', 30), ('Charlie', 20)]
    youngest = heapq.nsmallest(1, people, key=lambda x: x[1])
    print(youngest)  # Output: [('Charlie', 20)]
    

    Here, the key function “lambda x: x[1]” returns the second element of each tuple (i.e., the person’s age), which is used for comparison.

  • The heapq module can also be used to sort a list in ascending order using the heapq.heapreplace() function. This function pops the smallest element from the heap and replaces it with a new element, effectively sorting the list. For example:

    a = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
    heapq.heapify(a)
    sorted_list = [heapq.heappop(a) for _ in range(len(a))]
    

    In this example, the heapq.heappop() function is used in a list comprehension to pop each element from the heap and append it to a new list, resulting in a sorted list. However, note that this is not the most efficient way to sort a list in Python. The built-in sorted() function is faster and should be used instead.