Builds on Top-K Heap, which you haven’t finished yet.
You're given lists, an array holding k linked lists, where each individual list is already sorted from smallest to largest.
Combine every node from all k lists into a single linked list that is itself sorted from smallest to largest, and return its head.
K-Way Merge
K sorted lists ➔ min-heap of current heads ➔ O(N log K) merged stream| 1 | import heapq |
| 2 | from typing import Iterator, List, Optional, Tuple |
| 3 | |
| 4 | class ListNode: |
| 5 | def __init__(self, val: int = 0, next: Optional["ListNode"] = None): |
| 6 | self.val = val |
| 7 | self.next = next |
| 8 | |
| 9 | |
| 10 | def k_way_merge(sources: List[Iterator]) -> Iterator: |
| 11 | """K-Way Merge: every value of k sorted sources, in sorted order, holding one value per source. |
| 12 | Precondition: each source yields its values in ascending order and never yields None. |
| 13 | N = total values. K = number of sources. |
| 14 | Time: O(N log K) | Space: O(K) heap |
| 15 | """ |
| 16 | heap = [] # Min-heap of (value, source index); holds one unread value per source still open |
| 17 | |
| 18 | # 1. Seed One Candidate Per Source: |
| 19 | # Every source starts with its own first value, so the heap holds exactly one head per open source. |
| 20 | for i, source in enumerate(sources): |
| 21 | first = next(source, None) |
| 22 | if first is not None: |
| 23 | heapq.heappush(heap, (first, i)) # i breaks ties: values from two sources can be equal |
| 24 | |
| 25 | while heap: |
| 26 | # 2. Pop the Smallest Head and Yield It: |
| 27 | # heap[0] is always the smallest unread value across every open source, so popping it is safe to output now. |
| 28 | value, i = heapq.heappop(heap) # Smallest head across all sources |
| 29 | yield value |
| 30 | |
| 31 | # 3. Refill From the Same Source Only: |
| 32 | # A source advances only when its own head is used, which is what keeps every later pop in sorted order. |
| 33 | after = next(sources[i], None) |
| 34 | if after is not None: |
| 35 | heapq.heappush(heap, (after, i)) # Refill from the same source |
heap holds exactly one node per unfinished list, its current head, and every node already linked after dummy is no larger than any node in heap.The smallest remaining value overall is always one of the lists' current heads, so the heap only needs one node per list. Pop the smallest head, attach it after curr, and push that same list's node.next: a list advances only when its own head is used, which keeps the output sorted. Put the list index i in each tuple, so equal values are compared by index, never by ListNode.
See it run
K-Way Merge, step by step
heap []yielded 0
The first rows are the sources, with the position in the source as the column number. Below them are the heap value and source index at each heap array position, then the merged output.
in heapthe value of this source that is in the heap nowusedthe value popped in this stepyieldthe value written to the merged outputpushthe heap entry pushed in this steptiea heap entry whose value equals another entry's value
3 sorted sources: [[1, 4, 7, 9], [1, 3, 5], [2, 6, 8, 10]]. The heap is empty and nothing is merged yet.
Gotcha 1 of 4
- The trap
Gotcha 1: Pushing
(node.val, node)with no tie-breakerMistakeheapq.heappush(heap, (node.val, node))Fixheapq.heappush(heap, (node.val, i, node))What happens: Python compares tuples one item at a time. When two heads have the same
val, it compares the next items. Withouti, the next items are twoListNodeobjects.ListNodedefines no order, so Python raises a TypeError. Withiin the middle, equal values are ordered by list index, and the nodes are never compared.Example: Merge k Sorted Lists (LC 23): lists = [[1, 4, 5], [1, 3, 4], [2, 6]]Mistake: raises TypeErrorFix: returns [1, 1, 2, 3, 4, 4, 5, 6]
- Only the single value at rank k is asked, as in Kth Smallest Element in a Sorted Matrix (LC 378): nothing before it is used, so there is no merged output to build or read, only
k - 1pops to skip. See Best-First Heap
Deep Dive
Why one heap entry per source is enough
Top-K Heap already taught the heap. This card adds one rule about what the heap holds: exactly one entry for each source that still has values, and that entry is the source's current head.
Each source is sorted. So the smallest unread value of a source is always its head. The smallest unread value over all sources is therefore the smallest of the K heads. The heap root heap[0] is that value.
Pop the root and yield it. Now the source it came from has a new head, the value right after it. Push that value with the same source index i. Every other source keeps its head, so the heap again holds one entry per open source.
The output is sorted for this reason. Every value still unread is at least as large as some head in the heap, and the popped value was the smallest head. So no later pop can be smaller than an earlier one.
A source advances only when its own head is used. If the code advanced another source, one of its values would sit in the heap before its turn, and the heap would hold two entries from one source.
The merge traced on three sources
Take sources = [[1, 4, 7, 9], [1, 3, 5], [2, 6, 8, 10]]. There are K = 3 sources and N = 11 values.
The seed step pushes (1, 0), (1, 1) and (2, 2). The heap array is [(1, 0), (1, 1), (2, 2)]. The two entries with value 1 are equal in the first item, so Python compares the source index next, and (1, 0) comes first.
The first pop returns (1, 0) and yields 1. Source 0 moves on to 4, so (4, 0) is pushed. The next pop returns (1, 1), yields 1, and pushes (3, 1). The merged output starts [1, 1, ...].
Source 1 ends when (5, 1) is popped. It has no value after 5, so nothing is pushed and the heap shrinks to two entries. Source 0 ends at 9 and source 2 ends at 10 in the same way.
The merged output is [1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]. The heap never holds more than 3 entries, which is K. The LeetCode example of Merge k Sorted Lists (LC 23), [[1, 4, 5], [1, 3, 4], [2, 6]], gives [1, 1, 2, 3, 4, 4, 5, 6] in the same way.
Why the tuple needs the source index i
heapq compares tuples one item at a time. A tuple (value, node) is compared by value first. When two values are equal, Python compares the second items.
In Merge k Sorted Lists (LC 23) the second items would be two ListNode objects. ListNode defines no order, so Python raises TypeError: '<' not supported between instances of .... It happens as soon as two heads have equal values, as in [[1, 4, 5], [1, 3, 4], [2, 6]], where two lists start with 1.
The tuple (value, i, node) fixes this. Two entries with equal values are ordered by i. Two entries never share the same i, because the heap holds one entry per source. So Python never reaches the node item.
The index i is also how the code knows which source to refill after a pop. In the linked-list version, node.next gives the next head, and i is only a tie-breaker. In the iterator version, sources[i] is the source to read.
The cost, and why merging two at a time is slower
Let N be the total number of values and K the number of sources. Each value is pushed once and popped once. The heap holds at most K entries, so each push or pop costs O(log K). The total time is O(N log K). The extra space is O(K) for the heap.
Merging the sources two at a time in a row, the first with the second, then the result with the third, and so on, also gives the right answer. But the result grows, and the early values are copied again in every later merge.
For the example above, the row merges write 7 values and then 11 values, which is 18 writes, against 11 yields for the heap. The gap grows with K. In one test run with K = 100 sources of 100 random values each, the row merges wrote 504,900 values. The heap merge made 87,049 comparisons.
Merging the sources in pairs, like merge sort, also reaches O(N log K) time. It needs no heap. The heap version is the one that works on streams, because it reads each source one value at a time and never needs a whole source in memory.
60-Second Pitch
"Each source is sorted, so the smallest unread value overall is always the head of one source. I keep exactly those heads in a min-heap, one per source. I pop the smallest head and output it. Then I push the next value of the same source, because only that source has a new head. The output stays sorted, because nothing smaller can still be unread. I store (value, i) or (value, i, node), so equal values are ordered by the source index and Python never compares nodes. Each of the N values is pushed and popped once in a heap of at most K entries. The time is O(N log K) and the extra space is O(K). Merging two lists at a time would cost O(N * K)."
Common Questions
- Why does the heap hold only one entry per source?
A source is sorted, so its smallest unread value is its head. Only the heads can be the overall minimum. A value behind a head can be output only after the head, so it enters the heap when its head leaves.
- Why do you put
iin the tuple? Equal values would make Python compare the next tuple item. With
ListNodeobjects that raises aTypeError. With(value, i, node)the index orders the equal values, and the nodes are never compared.- What is the time and space complexity?
The time is
O(N log K), because each of the N values is pushed and popped once in a heap of at most K entries. The extra space is O(K). In the linked-list version the nodes are relinked, so no new nodes are created.- How does this differ from merging the lists two at a time?
Merging in a row copies the early values again in every later merge, so the time is O(N * K). Merging in pairs reaches
O(N log K)without a heap. The heap version also works when the sources are streams.- When would you pick this card over Top-K Heap?
Top-K Heap holds the k best items of one unsorted collection, and the heap size is k. This card holds one head for each sorted source, and the heap size is the number of sources. Use it when the inputs are already sorted.
SWE Reasoning
Spot it
- Several inputs are each sorted, and the answer needs their values in sorted order, as in Merge k Sorted Lists (LC 23).
- Only the smallest items across sorted sources are needed, so merging everything would do too much work.
- The sources are streams or files that must be read one value at a time.
Kis much smaller thanN, so a heap of sizeKis cheap.
The invariant, step by step
- Push the first value of each non-empty source as
(value, i). The heap now holds one head per open source. heap[0]is the smallest head. Each source is sorted, so it is also the smallest unread value of all sources.- Pop it and yield it, or link it after
curr. Everything already yielded is at most every entry in the heap. - Push the next value of the same source
i. If the source has no next value, push nothing. - Repeat until the heap is empty. The heap held one entry per open source at every step, so it never grew beyond
K.
Traps
- No tie-breaker in the tuple.
(node.val, node)compares twoListNodeobjects when two values are equal, and Python raises aTypeError. Use(node.val, i, node). - Pushing an empty list. An empty list is
None, andNone.valraises anAttributeError. Checkif node:before the first push. - Pushing
node.nextwithout a check. At the end of a listnode.nextisNone, andnode.next.valraises anAttributeError. Checkif node.next:. - Refilling from the wrong source. Pushing the next value of another source puts two entries of one source in the heap. The output can then be out of order.
- Merging the lists two at a time in a row. The result is right, but the time is O(N * K) instead of
O(N log K).
Practice
Builds on Top-K Heap, which you haven’t finished yet.
You're given lists, an array holding k linked lists, where each individual list is already sorted from smallest to largest.
Combine every node from all k lists into a single linked list that is itself sorted from smallest to largest, and return its head.
K-Way Merge
K sorted lists ➔ min-heap of current heads ➔ O(N log K) merged stream| 1 | import heapq |
| 2 | from typing import Iterator, List, Optional, Tuple |
| 3 | |
| 4 | class ListNode: |
| 5 | def __init__(self, val: int = 0, next: Optional["ListNode"] = None): |
| 6 | self.val = val |
| 7 | self.next = next |
| 8 | |
| 9 | |
| 10 | def k_way_merge(sources: List[Iterator]) -> Iterator: |
| 11 | """K-Way Merge: every value of k sorted sources, in sorted order, holding one value per source. |
| 12 | Precondition: each source yields its values in ascending order and never yields None. |
| 13 | N = total values. K = number of sources. |
| 14 | Time: O(N log K) | Space: O(K) heap |
| 15 | """ |
| 16 | heap = [] # Min-heap of (value, source index); holds one unread value per source still open |
| 17 | |
| 18 | # 1. Seed One Candidate Per Source: |
| 19 | # Every source starts with its own first value, so the heap holds exactly one head per open source. |
| 20 | for i, source in enumerate(sources): |
| 21 | first = next(source, None) |
| 22 | if first is not None: |
| 23 | heapq.heappush(heap, (first, i)) # i breaks ties: values from two sources can be equal |
| 24 | |
| 25 | while heap: |
| 26 | # 2. Pop the Smallest Head and Yield It: |
| 27 | # heap[0] is always the smallest unread value across every open source, so popping it is safe to output now. |
| 28 | value, i = heapq.heappop(heap) # Smallest head across all sources |
| 29 | yield value |
| 30 | |
| 31 | # 3. Refill From the Same Source Only: |
| 32 | # A source advances only when its own head is used, which is what keeps every later pop in sorted order. |
| 33 | after = next(sources[i], None) |
| 34 | if after is not None: |
| 35 | heapq.heappush(heap, (after, i)) # Refill from the same source |
heap holds exactly one node per unfinished list, its current head, and every node already linked after dummy is no larger than any node in heap.The smallest remaining value overall is always one of the lists' current heads, so the heap only needs one node per list. Pop the smallest head, attach it after curr, and push that same list's node.next: a list advances only when its own head is used, which keeps the output sorted. Put the list index i in each tuple, so equal values are compared by index, never by ListNode.
See it run
K-Way Merge, step by step
heap []yielded 0
The first rows are the sources, with the position in the source as the column number. Below them are the heap value and source index at each heap array position, then the merged output.
in heapthe value of this source that is in the heap nowusedthe value popped in this stepyieldthe value written to the merged outputpushthe heap entry pushed in this steptiea heap entry whose value equals another entry's value
3 sorted sources: [[1, 4, 7, 9], [1, 3, 5], [2, 6, 8, 10]]. The heap is empty and nothing is merged yet.
Gotcha 1 of 4
- The trap
Gotcha 1: Pushing
(node.val, node)with no tie-breakerMistakeheapq.heappush(heap, (node.val, node))Fixheapq.heappush(heap, (node.val, i, node))What happens: Python compares tuples one item at a time. When two heads have the same
val, it compares the next items. Withouti, the next items are twoListNodeobjects.ListNodedefines no order, so Python raises a TypeError. Withiin the middle, equal values are ordered by list index, and the nodes are never compared.Example: Merge k Sorted Lists (LC 23): lists = [[1, 4, 5], [1, 3, 4], [2, 6]]Mistake: raises TypeErrorFix: returns [1, 1, 2, 3, 4, 4, 5, 6]
- Only the single value at rank k is asked, as in Kth Smallest Element in a Sorted Matrix (LC 378): nothing before it is used, so there is no merged output to build or read, only
k - 1pops to skip. See Best-First Heap
Deep Dive
Why one heap entry per source is enough
Top-K Heap already taught the heap. This card adds one rule about what the heap holds: exactly one entry for each source that still has values, and that entry is the source's current head.
Each source is sorted. So the smallest unread value of a source is always its head. The smallest unread value over all sources is therefore the smallest of the K heads. The heap root heap[0] is that value.
Pop the root and yield it. Now the source it came from has a new head, the value right after it. Push that value with the same source index i. Every other source keeps its head, so the heap again holds one entry per open source.
The output is sorted for this reason. Every value still unread is at least as large as some head in the heap, and the popped value was the smallest head. So no later pop can be smaller than an earlier one.
A source advances only when its own head is used. If the code advanced another source, one of its values would sit in the heap before its turn, and the heap would hold two entries from one source.
The merge traced on three sources
Take sources = [[1, 4, 7, 9], [1, 3, 5], [2, 6, 8, 10]]. There are K = 3 sources and N = 11 values.
The seed step pushes (1, 0), (1, 1) and (2, 2). The heap array is [(1, 0), (1, 1), (2, 2)]. The two entries with value 1 are equal in the first item, so Python compares the source index next, and (1, 0) comes first.
The first pop returns (1, 0) and yields 1. Source 0 moves on to 4, so (4, 0) is pushed. The next pop returns (1, 1), yields 1, and pushes (3, 1). The merged output starts [1, 1, ...].
Source 1 ends when (5, 1) is popped. It has no value after 5, so nothing is pushed and the heap shrinks to two entries. Source 0 ends at 9 and source 2 ends at 10 in the same way.
The merged output is [1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]. The heap never holds more than 3 entries, which is K. The LeetCode example of Merge k Sorted Lists (LC 23), [[1, 4, 5], [1, 3, 4], [2, 6]], gives [1, 1, 2, 3, 4, 4, 5, 6] in the same way.
Why the tuple needs the source index i
heapq compares tuples one item at a time. A tuple (value, node) is compared by value first. When two values are equal, Python compares the second items.
In Merge k Sorted Lists (LC 23) the second items would be two ListNode objects. ListNode defines no order, so Python raises TypeError: '<' not supported between instances of .... It happens as soon as two heads have equal values, as in [[1, 4, 5], [1, 3, 4], [2, 6]], where two lists start with 1.
The tuple (value, i, node) fixes this. Two entries with equal values are ordered by i. Two entries never share the same i, because the heap holds one entry per source. So Python never reaches the node item.
The index i is also how the code knows which source to refill after a pop. In the linked-list version, node.next gives the next head, and i is only a tie-breaker. In the iterator version, sources[i] is the source to read.
The cost, and why merging two at a time is slower
Let N be the total number of values and K the number of sources. Each value is pushed once and popped once. The heap holds at most K entries, so each push or pop costs O(log K). The total time is O(N log K). The extra space is O(K) for the heap.
Merging the sources two at a time in a row, the first with the second, then the result with the third, and so on, also gives the right answer. But the result grows, and the early values are copied again in every later merge.
For the example above, the row merges write 7 values and then 11 values, which is 18 writes, against 11 yields for the heap. The gap grows with K. In one test run with K = 100 sources of 100 random values each, the row merges wrote 504,900 values. The heap merge made 87,049 comparisons.
Merging the sources in pairs, like merge sort, also reaches O(N log K) time. It needs no heap. The heap version is the one that works on streams, because it reads each source one value at a time and never needs a whole source in memory.
60-Second Pitch
"Each source is sorted, so the smallest unread value overall is always the head of one source. I keep exactly those heads in a min-heap, one per source. I pop the smallest head and output it. Then I push the next value of the same source, because only that source has a new head. The output stays sorted, because nothing smaller can still be unread. I store (value, i) or (value, i, node), so equal values are ordered by the source index and Python never compares nodes. Each of the N values is pushed and popped once in a heap of at most K entries. The time is O(N log K) and the extra space is O(K). Merging two lists at a time would cost O(N * K)."
Common Questions
- Why does the heap hold only one entry per source?
A source is sorted, so its smallest unread value is its head. Only the heads can be the overall minimum. A value behind a head can be output only after the head, so it enters the heap when its head leaves.
- Why do you put
iin the tuple? Equal values would make Python compare the next tuple item. With
ListNodeobjects that raises aTypeError. With(value, i, node)the index orders the equal values, and the nodes are never compared.- What is the time and space complexity?
The time is
O(N log K), because each of the N values is pushed and popped once in a heap of at most K entries. The extra space is O(K). In the linked-list version the nodes are relinked, so no new nodes are created.- How does this differ from merging the lists two at a time?
Merging in a row copies the early values again in every later merge, so the time is O(N * K). Merging in pairs reaches
O(N log K)without a heap. The heap version also works when the sources are streams.- When would you pick this card over Top-K Heap?
Top-K Heap holds the k best items of one unsorted collection, and the heap size is k. This card holds one head for each sorted source, and the heap size is the number of sources. Use it when the inputs are already sorted.
SWE Reasoning
Spot it
- Several inputs are each sorted, and the answer needs their values in sorted order, as in Merge k Sorted Lists (LC 23).
- Only the smallest items across sorted sources are needed, so merging everything would do too much work.
- The sources are streams or files that must be read one value at a time.
Kis much smaller thanN, so a heap of sizeKis cheap.
The invariant, step by step
- Push the first value of each non-empty source as
(value, i). The heap now holds one head per open source. heap[0]is the smallest head. Each source is sorted, so it is also the smallest unread value of all sources.- Pop it and yield it, or link it after
curr. Everything already yielded is at most every entry in the heap. - Push the next value of the same source
i. If the source has no next value, push nothing. - Repeat until the heap is empty. The heap held one entry per open source at every step, so it never grew beyond
K.
Traps
- No tie-breaker in the tuple.
(node.val, node)compares twoListNodeobjects when two values are equal, and Python raises aTypeError. Use(node.val, i, node). - Pushing an empty list. An empty list is
None, andNone.valraises anAttributeError. Checkif node:before the first push. - Pushing
node.nextwithout a check. At the end of a listnode.nextisNone, andnode.next.valraises anAttributeError. Checkif node.next:. - Refilling from the wrong source. Pushing the next value of another source puts two entries of one source in the heap. The output can then be out of order.
- Merging the lists two at a time in a row. The result is right, but the time is O(N * K) instead of
O(N log K).