Функция для нахождения медианы в потоке чисел

import heapq

class MedianFinder:
    def __init__(self):
        # Мин-куча для хранения наибольших значений
        self.min_heap = []
        # Макс-куча для хранения наименьших значений
        self.max_heap = []

    def addNum(self, num):
        # Добавляем в макс-кучу
        heapq.heappush(self.max_heap, -num)

        # Балансировка куч
        if (self.max_heap and self.min_heap and
            (-self.max_heap[0]) > self.min_heap[0]):
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))

        # Поддержание баланса размеров куч
        if len(self.max_heap) > len(self.min_heap) + 1:
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        if len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

    def findMedian(self):
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2

# Пример использования:
mf = MedianFinder()
mf.addNum(1)
mf.addNum(2)
print(mf.findMedian()) # Выводит 1.5
mf.addNum(3)
print(mf.findMedian()) # Выводит 2