Let's solve an interesting problem - merging
Let’s break down the problem statement: We have
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.
The simplest approach that comes to mind is the
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 n
is the total number of elements in all sublists, and its space complexity is
Let's now look for an optimized approach.
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.
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 n
is the total number of elements and k
is the number of lists and the space complexity of the code is n
is the total number of elements and m
is the number of lists.
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