How to merge k sorted lists in Python

Let's solve an interesting problem - merging kk sorted lists into one sorted list. This problem is a common one that you might encounter during coding interviews or when dealing with real-world datasets.

1 of 12

Problem statement

Let’s break down the problem statement: We have kk sorted lists (where k can be any non-negative integer), and we want to merge all these lists into one sorted list.

Here is a simple example to illustrate the problem:

Input: [[1, 4, 5], [1, 3, 4], [2, 6]]
Output: [1, 1, 2, 3, 4, 4, 5, 6]

Let's move on to solving the problem.

Brute force approach

The simplest approach that comes to mind is the brute force methodBrute force approach involves checking every possible solution without optimizing or applying any specific algorithmic technique., where we add all elements from all lists into a new list and sort it. Here’s how we do it in Python:

def merge_lists_bruteforce(lists):
    result = []
    for l in lists:
        result.extend(l)
    return sorted(result)
print(merge_lists_bruteforce([[1, 4, 5], [1, 3, 4], [2, 6]]))

While this code works fine and is straightforward, it’s not the most efficient solution. Its time complexity is O(nlogn)O(n log n) because we are sorting the entire list, where n is the total number of elements in all sublists, and its space complexity is O(n)O(n), where n is the total number of elements in all sublists.

Optimized approach

Let's now look for an optimized approach.

Using heap

We can use a data structure known as a heap to solve this problem more efficiently. A heap is a binary tree-like data structure that maintains the heap property. The key at each parent node is either less than or equal to its child nodes (in a min-heap) or more than or equal to its child nodes (in a max heap).

To explore the difference between min heap and max heap, click here.

We will use a min-heap to solve this problem as we want the smallest elements to be popped first. The Python heapq module provides an implementation of heaps via priority queues. The elements of the heap can be any ordered pairs, and the pair is sorted by the first element, then the second, and so on.

Example

Let’s understand it with an example:

import heapq  # We import the heapq module. This module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm.
def merge_lists(lists):
    # For each list, we add its first element to our heap along with the list's index and the element's index within that list.
    # We only add the first element of lists that are not empty to prevent IndexError.
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    
    heapq.heapify(heap)  # This transforms list x into a heap, in-place, in linear time.
    result = []  # This will hold our final merged list.
    
    while heap:  # While our heap is not empty.
        val, list_ind, element_ind = heapq.heappop(heap)  # We pop (remove and return) the smallest element from the heap, along with its list index and its element index within that list.
        result.append(val)  # We append this smallest element to our result list.
        
        # If there are more elements in the list from which the smallest element came, we push the next element to the heap.
        # We only push the next element if it exists to prevent IndexError.
        if element_ind + 1 < len(lists[list_ind]):
            next_tuple = (lists[list_ind][element_ind + 1], list_ind, element_ind + 1)  # The next element along with the list index and its element index.
            heapq.heappush(heap, next_tuple)  # We push the next_tuple into the heap.
            
    return result  # We return our final merged list.
print(merge_lists([[1, 4, 5], [1, 3, 4], [2, 6]]))  # Testing our function.

In the code above, we first initialize a heap with the first element of each list along with the list index and element index. We then pop the heap to get the smallest element and add it to the result list. If there are remaining elements in that list, we push the next element into the heap. We continue this process until the heap is empty.

This solution has a time complexity O(nlogk)O(n log k), where n is the total number of elements and k is the number of lists and the space complexity of the code is O(n+m)O(n + m), where n is the total number of elements and m is the number of lists.

Conclusion

Merging k sorted lists is a common problem in both coding interviews and real-world scenarios. While the brute force approach is easier to understand, the heap-based approach is more efficient and faster for large inputs. Python’s built-in heapq module provides a convenient way to solve this problem using a min-heap.

Free Resources

Copyright ©2025 Educative, Inc. All rights reserved