merge ordered lists in python. compare heap merge vs naive implementation

merge_lists.py

import heapq
import operator
import functools
import time
from random import randint


def timeit(func):
    @functools.wraps(func)
    def newfunc(*args, **kwargs):
        startTime = time.time()
        ret = func(*args, **kwargs)
        elapsedTime = time.time() - startTime
        print('function [{}] finished in {} ms'.format(
            func.__name__, elapsedTime * 1000.))
        return ret
    return newfunc


def generate_sorted_lists(num_of_lists, num_of_elements):
    result = []
    for n in range(num_of_lists):
        start = randint(0, 1000)
        result.append(list(range(start, start + num_of_elements)))
    return result


@timeit
def merge_sorted_lists_with_heap(lists):
    merged_list = []
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)
    while heap:
        val, list_ind, element_ind = heapq.heappop(heap)
        merged_list.append(val)
        if element_ind + 1 < len(lists[list_ind]):
            next_tuple = (lists[list_ind][element_ind + 1],
                          list_ind,
                          element_ind + 1)
            heapq.heappush(heap, next_tuple)
    return merged_list


@timeit
def merge_sorted_lists(lists):
    l = len(lists[0])
    result = []
    n = len(lists)
    indexes = [0] * n
    for j in range(l * n):
        templist = []
        for i, index in enumerate(indexes):
            if index < l:
                templist.append(lists[i][index])
            else:
                templist.append(9999999999999999)
        minindex, minvalue = min(enumerate(templist), key=operator.itemgetter(1))
        indexes[minindex] = indexes[minindex] + 1
        result.append(minvalue)
    return result


def test(num_of_lists, length_of_list):
    print('merging %i lists of lenght %i' % (num_of_lists, length_of_list))
    lists = generate_sorted_lists(num_of_lists, length_of_list)
    a = merge_sorted_lists_with_heap(lists)
    b = merge_sorted_lists(lists)
    assert(a == b)
    print('===================================================================')


def run_tests():
    test(10, 100)
    test(100, 10)
    test(100, 100)
    test(100, 1000)
    test(2, 100000)
    test(3, 1000000)


run_tests()



merging 10 lists of lenght 100
function [merge_sorted_lists_with_heap] finished in 0.8451938629150391 ms
function [merge_sorted_lists] finished in 4.227876663208008 ms
===================================================================
merging 100 lists of lenght 10
function [merge_sorted_lists_with_heap] finished in 1.0259151458740234 ms
function [merge_sorted_lists] finished in 27.924060821533203 ms
===================================================================
merging 100 lists of lenght 100
function [merge_sorted_lists_with_heap] finished in 11.030912399291992 ms
function [merge_sorted_lists] finished in 288.85793685913086 ms
===================================================================
merging 100 lists of lenght 1000
function [merge_sorted_lists_with_heap] finished in 113.91901969909668 ms
function [merge_sorted_lists] finished in 3084.4509601593018 ms
===================================================================
merging 2 lists of lenght 100000
function [merge_sorted_lists_with_heap] finished in 163.93303871154785 ms
function [merge_sorted_lists] finished in 395.4272270202637 ms
===================================================================
merging 3 lists of lenght 1000000
function [merge_sorted_lists_with_heap] finished in 2540.393114089966 ms
function [merge_sorted_lists] finished in 6837.332010269165 ms
===================================================================