Merge k sorted lists into a sorted list in O(n lg k), HINT: use a min-heap, *CORMEN Introduction to Algorithms exercise

k-list-heap-merge.py

from math import floor

class Heap:
    def __init__(self):
        self.nodes = list()
        self.heap_size = -1

    def __left(self, i):
        return i * 2 + 1

    def __right(self, i):
        return i * 2 + 2

    def __parent(self, i):
        return int(floor(i/2))

    def insert(self, element):
        element = element[:] # original lists shouldn't be edited.
        self.nodes.append(element)
        self.heap_size += 1
        i = self.heap_size
        while i > 0 and self.nodes[self.__parent(i)][0] > self.nodes[i][0]:
            self.nodes[i], self.nodes[self.__parent(i)] = self.nodes[self.__parent(i)], self.nodes[i]
            i = self.__parent(i)

    def heapify(self, i):
        min_i = i
        if self.heap_size >= self.__left(i):
            if self.nodes[i][0] > self.nodes[self.__left(i)][0]:
                min_i = self.__left(i)
        if self.heap_size >= self.__right(i):
            if self.nodes[min_i][0] > self.nodes[self.__right(i)][0]:
                min_i = self.__right(i)
        if min_i != i:
            self.nodes[i], self.nodes[min_i] = self.nodes[min_i], self.nodes[i]
            self.heapify(min_i)

    def extract(self):
        if self.heap_size >= 0:
            x = self.nodes[0][0]
            del self.nodes[0][0]
            if not len(self.nodes[0]):
                self.nodes[self.heap_size - 1], self.nodes[0] = self.nodes[0], self.nodes[self.heap_size - 1]
                del self.nodes[self.heap_size - 1]
                self.heap_size -= 1
            if self.heap_size >= 0:
                self.heapify(0)
            return x

    def extract_all(self):
        data = list()
        while self.heap_size >= 0:
            data.append(self.extract())
        return data
   
heap = Heap()
heap.insert([2, 4, 5, 8, 13, 22])
heap.insert([1, 3, 6, 10, 11, 13])
heap.insert([4, 5, 7, 8, 9, 40])
heap.insert([5, 6, 10, 12, 20, 22])

print heap.extract_all()