Merge K sorted array

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)

view raw merge_sorted_array_1.py hosted with ❤ by GitHub

	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

view raw merge_sorted_array_2.py hosted with ❤ by GitHub

	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

view raw merged_sorted_array_3.py hosted with ❤ by GitHub