Anonymous Anonymous - 6 months ago 183 0

No description

Python

Merge k lists

__author__ = 'ssaraswati'

from heapq import heappush, heappop


def combine(inputList):
    listsCount =  len(inputList)
    outputList = []
    heap = []
    if listsCount == 1:
        return inputList[0]
    i = 0

    while i < listsCount:
        heappush(heap, [inputList[i].pop(0), i])
        i += 1

    while len(heap) != 0:
        nextVal = heappop(heap)
        outputList.append(nextVal[0])
        if len(inputList[nextVal[1]]) > 0:
            heappush(heap, [inputList[nextVal[1]].pop(0), nextVal[1]])
    return outputList

def test1():
    result4 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 20]
    input4 = [result4]
    print "****Test 1****"
    assert combine(input4) == result4, "not sorting 1 list correctly"


def test2():
    listA = [1, 3, 5, 7]
    listB = [2, 4, 6, 8, 10]
    input2 = [listA, listB]
    result2 = [1, 2, 3, 4, 5, 6, 7, 8, 10]
    print "****Test 2****"
    assert combine(input2) == result2, "not sorting 2 lists correctly"


def test3():
    listA = [1, 3, 5, 7]
    listB = [2, 4, 6, 8, 10]
    listC = [0, 20]
    input3 = [listA, listB, listC]
    result3 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 20]

    print "****Test 3****"
    assert combine(input3) == result3, "not sorting 3 lists correctly"

def testRunner():
    test1();
    test2();
    test3();

testRunner()
Comments