Heap
Merge K Sorted Lists
DESCRIPTION (inspired by Leetcode.com)
Given k linked lists, each sorted in ascending order, in a list lists, write a function to merge the input lists into one sorted linked list.
Example 1:
Inputs:
lists = [[3,4,6],[2,3,5],[-1,6]] 3 -> 4 -> 6 2 -> 3 -> 5 -1 -> 6
Output:
[-1,2,3,3,4,5,6,6]
-1 -> 2 -> 3 -> 3 -> 4 -> 5 -> 6 -> 6
Explanation
Step 1: Initialize the heap
- The element's value
- The array's index (which array it came from)
- The element's index within that array (starting at 0)
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
initialize heap
0 / 3
Step 2: Build the merged array
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
add -1 to heap
0 / 1
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
initialize result array
0 / 1
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
add -1 to result
0 / 1
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
add 6 to result
0 / 1
Solution
from heapq import heappush, heappopdef mergeKLists(lists):if not lists:return []has_non_empty = any(lst for lst in lists)if not has_non_empty:return []heap = []for i, lst in enumerate(lists):if lst:heappush(heap, (lst[0], i, 0))result = []while heap:value, list_idx, elem_idx = heappop(heap)result.append(value)if elem_idx + 1 < len(lists[list_idx]):next_value = lists[list_idx][elem_idx + 1]heappush(heap, (next_value, list_idx, elem_idx + 1))return result
merge k sorted arrays
0 / 19
Test Your Knowledge
Login to take the complexity quiz and track your progress
Complexity Analysis
Time Complexity: O(n * log k) where n is the total number of elements across all input arrays and k is the number of arrays. We iterate over all elements from all arrays. At each iteration, we push and pop elements from the heap, which takes O(log k) time.
Space Complexity: O(k) where k is the number of arrays. We use a min-heap of size k to store the next unmerged element from each of the k arrays.
Mark as read
Unlock Premium Coding Content
On This Page
Your account is free and you can post anonymously if you choose.