Merge K sorted Arrays using Python

#Create a node to save in the minHeap being created
class HeapNode(object):
    def __init__(self):
        self.array = 0
        self.data = 0
        self.index = -1

#min Heap implementation
class BinHeap:
    def __init__(self):
        self.heapList = [0]
        self.currentSize = 0

    def isEmpty(self):
        return self.currentSize == 0

    def percUp(self, i):
        while i // 2 > 0:
            if self.heapList[i].data < self.heapList[i // 2].data:
                tmp = self.heapList[i // 2]
                self.heapList[i // 2] = self.heapList[i]
                self.heapList[i] = tmp
            i = i // 2

    def insert(self, k):
        self.heapList.append(k)
        self.currentSize = self.currentSize + 1
        self.percUp(self.currentSize)

    def percDown(self, i):
        while (i * 2) <= self.currentSize:
            mc = self.minChild(i)
            if self.heapList[i].data > self.heapList[mc].data:
                tmp = self.heapList[i]
                self.heapList[i] = self.heapList[mc]
                self.heapList[mc] = tmp
            i = mc

    def minChild(self, i):

        if i * 2 + 1 > self.currentSize:
            return i * 2
        else:
            if self.heapList[i * 2].data < self.heapList[i * 2 + 1].data:
                return i * 2
     W       else:
                return i * 2 + 1

    def delMin(self):
        retval = self.heapList[1]
        self.heapList[1] = self.heapList[self.currentSize]
        self.currentSize = self.currentSize - 1
        self.heapList.pop()
        self.percDown(1)
        return retval

    def buildHeap(self, alist):
        i = len(alist) // 2
        self.currentSize = len(alist)
        self.heapList = [0] + alist[:]
        while (i > 0):
            self.percDown(i)
            i = i - 1

#merge k sorted arrays function

def merge(arr):
    pq = BinHeap()
    size = 0

    for i in range(len(arr)):
        size +=len(arr[i])
        if len(arr[i]) > 0:
            node = HeapNode()
            node.array = i
            node.data = arr[i][0]
            node.index = 0
            pq.insert(node)
    result = [0]*size
    i = 0
    while not pq.isEmpty():
        n = pq.delMin()
        result[i] = n.data
        newIndex = n.index +1
        if newIndex < len(arr[n.array]):
            node = HeapNode()
            node.array = n.array
            node.data = arr[n.array][newIndex]
            node.index = newIndex
            pq.insert(node)
        i +=1
    return result


print(merge([[1,2,7],[3,4,6]]))
Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s