min waiting time problem

test.py

import heapq

class MyHeap(object):
    def __init__(self, initial=None, key=lambda x:x):
        self.key = key
        if initial:
            self._data = [(key(item), item) for item in initial]
            heapq.heapify(self._data)
        else:
            self._data = []

    def push(self, item):
        heapq.heappush(self._data, (self.key(item), item))

    def pop(self):
        return heapq.heappop(self._data)[1]

    def __len__(self):
        return len(self._data)


def solution(lst):
    tLst = sorted(lst)
    h = MyHeap(key=lambda x: x[1])
    h.push(tLst[0])
    current_time = 0
    index = 1
    total_wait_time = 0
    while len(h) > 0:
        next_served = h.pop()
        total_wait_time += next_served[1] * (1 + len(h))
        print("serving: ", next_served)
        next_time = current_time + next_served[1]
        while True:
            if index < len(tLst) and tLst[index][0] <= next_time:
                h.push(tLst[index])
                total_wait_time = total_wait_time + (next_time - tLst[index][0])
                index = index + 1
            else:
                break
        current_time = next_time
    print("average time: ", total_wait_time / len(lst))

solution([(0, 3), (1, 9), (2, 5)])