K-way merge sort

k_way_merge.py

from heapq import heappush, heappop

class StreamScore(object):

    def __init__(self, stream, key, value):
        self.stream = stream
        self.key = key
        self.value = value

    def __lt__(self, other):
        if self.key == other.key:
            return self.value < other.value
        return self.key < other.key

    def __le__(self, other):
        if self.key == other.key:
            return self.value <= other.value
        return self.key <= other.key


    def __eq__(self, other):
        return self.key == other.key

    def __ne__(self, other):
        return self.key != other.key

    def __gt__(self, other):
        if self.key == other.key:
            return self.value > other.value
        return self.key > other.key

    def __ge__(self, other):
        if self.key == other.key:
            return self.value >= other.value
        return self.key >= other.key
        
def k_way_merge(streams):
    sorted = []
    heap = []
    empty = [False]*len(streams)


    for i in range(len(streams)):
        if len(streams[i]) != 0:
            key, value = streams[i].pop(0)
        heappush(heap, StreamScore(i, key, value))

    while (not all(empty)):
        next = heappop(heap)

        sorted.append((next.key, next.value))

        if len(streams[next.stream]) == 0:
            empty[next.stream] = True
        else:
            key, value = streams[next.stream].pop(0)
            heappush(heap, StreamScore(next.stream, key, value))

    print(len(heap))

    return sorted

test = [
    [(2,100), (5, 200), (1, 50), (2, 200), (7, 1)],
    [(18, 50), (4, 66), (90, 1), (20, 300)],
    [(78, 3), (3, 78), (10, 2), (2, 100)]
]


sorted_test = [sorted(i) for i in test]

print(k_way_merge(sorted_test))