-
Notifications
You must be signed in to change notification settings - Fork 4.7k
Expand file tree
/
Copy pathmerge_sorted_k_lists.py
More file actions
62 lines (47 loc) · 1.48 KB
/
merge_sorted_k_lists.py
File metadata and controls
62 lines (47 loc) · 1.48 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
"""
Merge K Sorted Linked Lists
Merge k sorted linked lists into one sorted linked list using a heap
for efficient minimum extraction.
Reference: https://leetcode.com/problems/merge-k-sorted-lists/
Complexity:
Time: O(n log k) where n is total elements and k is number of lists
Space: O(k)
"""
from __future__ import annotations
from heapq import heapify, heappop, heapreplace
class ListNode:
"""Singly linked list node.
Args:
val: The node value.
"""
def __init__(self, val: int) -> None:
self.val = val
self.next: ListNode | None = None
def merge_k_lists(lists: list[ListNode | None]) -> ListNode | None:
"""Merge k sorted linked lists into a single sorted linked list.
Args:
lists: A list of head nodes of sorted linked lists.
Returns:
Head of the merged sorted linked list, or None if all are empty.
Examples:
>>> n1 = ListNode(1)
>>> n2 = ListNode(2)
>>> result = merge_k_lists([n1, n2])
>>> result.val
1
"""
dummy = node = ListNode(0)
heap: list[tuple[int, int, ListNode]] = []
for idx, head in enumerate(lists):
if head:
heap.append((head.val, idx, head))
heapify(heap)
while heap:
val, idx, n_val = heap[0]
if n_val.next is None:
heappop(heap)
else:
heapreplace(heap, (n_val.next.val, idx, n_val.next))
node.next = n_val
node = node.next
return dummy.next