This is LeetCode problem 23: merge K sorted array.
The solution is to use heap to merge sub array.
I provide three kinds of solution. Use the second solution to pass LeetCode.
import heapq | |
def merge_arr(arr): | |
result = [] | |
if len(arr) == 0: | |
return result | |
tmp = [] # tmp heap | |
# push the first item of all subarray into heap | |
# each heap item is [value, array index, the index inside subarray] | |
for i in range(len(arr)): | |
heapq.heappush(tmp, [arr[i][0], i, 0]) | |
while len(tmp) > 0: | |
# pop heap | |
next_item = heapq.heappop(tmp) | |
if next_item is None: | |
break | |
result.append(next_item[0]) | |
sub_arr = next_item[1] | |
next_index = next_item[2] | |
# add subarray next item | |
if next_index < len(arr[sub_arr]) - 1: | |
heapq.heappush(tmp, [arr[sub_arr][next_index + 1], sub_arr, next_index + 1]) | |
return result | |
if __name__ == "__main__": | |
arr =[[1, 6, 12 ], | |
[1, 9 ], | |
[23, 34, 90, 2000 ]] | |
print("arr:", arr) | |
result = merge_arr(arr) | |
print("merged:", result) |
import heapq | |
from typing import List | |
# Definition for singly-linked list. | |
# Leet code problem 23 | |
class ListNode: | |
def __init__(self, val=0, next=None): | |
self.val = val | |
self.next = next | |
class Solution: | |
def mergeKLists(self, lists: List[ListNode]) -> ListNode: | |
result = ListNode() | |
next_node = None | |
if len(lists) == 0: | |
return ListNode(val='') | |
tmp = [] # heap | |
pointer = [None] * len(lists) # keep trace each list position | |
for i in range(len(lists)): | |
if lists[i] is not None: | |
pointer[i] = lists[i] | |
heapq.heappush(tmp, [lists[i].val, i]) | |
if len(tmp) == 0: | |
return ListNode(val='') | |
count = 0 | |
while len(tmp) > 0: | |
tmp_node = heapq.heappop(tmp) | |
if tmp_node is not None: | |
array_index = tmp_node[1] | |
next_node = pointer[array_index] | |
# The head of list | |
if count == 0: | |
result = next_node | |
else: | |
prev_node.next = next_node | |
count += 1 | |
prev_node = next_node | |
next_node = next_node.next | |
pointer[array_index] = next_node | |
# push next node if not empty | |
if next_node is not None: | |
heapq.heappush(tmp, [next_node.val, array_index]) | |
return result | |
if __name__ == "__main__": | |
list1 = ListNode(val=1) | |
list1.next = ListNode(val=6) | |
list1.next.next = ListNode(val=12) | |
list2 = ListNode(val=1) | |
list2.next = ListNode(val=9) | |
list3 = ListNode(val=23) | |
list3.next = ListNode(val=34) | |
list3.next.next = ListNode(val=90) | |
list3.next.next.next = ListNode(val=2000) | |
lists =[ list1, list2, list3] | |
#lists = [] | |
solution = Solution() | |
sort_list = solution.mergeKLists(lists) | |
node = sort_list | |
print("merged:") | |
while node: | |
print(node.val, end=' ') | |
node = node.next |
import heapq | |
from typing import List | |
# Definition for singly-linked list. | |
class ListNode: | |
def __init__(self, val=0, next=None): | |
self.val = val | |
self.next = next | |
def __lt__(self, other): | |
return self.val < other.val | |
class Solution: | |
def mergeKLists(self, lists: List[ListNode]) -> ListNode: | |
result = ListNode() | |
next_node = None | |
if len(lists) == 0: | |
return ListNode(val='') | |
tmp = [] # heap | |
for i in range(len(lists)): | |
if lists[i] is not None: | |
heapq.heappush(tmp, lists[i]) | |
if len(tmp) == 0: | |
return ListNode(val='') | |
count = 0 | |
while len(tmp) > 0: | |
next_node = heapq.heappop(tmp) | |
if next_node is not None: | |
# The head of list | |
if count == 0: | |
result = next_node | |
else: | |
prev_node.next = next_node | |
count += 1 | |
prev_node = next_node | |
next_node = next_node.next | |
# push next node if not empty | |
if next_node is not None: | |
heapq.heappush(tmp, next_node) | |
return result | |
if __name__ == "__main__": | |
list1 = ListNode(val=1) | |
list1.next = ListNode(val=6) | |
list1.next.next = ListNode(val=12) | |
list2 = ListNode(val=1) | |
list2.next = ListNode(val=9) | |
list3 = ListNode(val=23) | |
list3.next = ListNode(val=34) | |
list3.next.next = ListNode(val=90) | |
list3.next.next.next = ListNode(val=2000) | |
lists =[ list1, list2, list3] | |
solution = Solution() | |
sort_list = solution.mergeKLists(lists) | |
node = sort_list | |
print("merged:") | |
while node: | |
print(node.val, end=' ') | |
node = node.next |