Question

How to merge N sorted lists in Python

Answer 1

import heapq

def merge_n_sorted_lists(lists):
    min_heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(min_heap, (lst[0], i, 0))

    merged_list = []
    while min_heap:
        val, list_idx, element_idx = heapq.heappop(min_heap)
        merged_list.append(val)
        if element_idx + 1 < len(lists[list_idx]):
            next_tuple = (lists[list_idx][element_idx + 1], list_idx, element_idx + 1)
            heapq.heappush(min_heap, next_tuple)

    return merged_list

lists = [
    [1, 2, 6, 8],
    [1, 2, 9],
    [2, 7]
]

print(merge_n_sorted_lists(lists))



'''
run:

[1, 1, 2, 2, 2, 6, 7, 8, 9]

'''

 

Learn SQL from A to Z at LearnSQL.com