implementation of max heap and min heap to solve running median problem

running_median.py

import sys


def left(i):
    return i * 2


def right(i):
    return i * 2 + 1


def parent(i):
    return i // 2


def max_heapify(heap, heap_size, i):
    l = left(i)
    r = right(i)

    if l <= heap_size and heap[l] > heap[i]:
        largest = l
    else:
        largest = i

    if r <= heap_size and heap[r] > heap[largest]:
        largest = r

    if i != largest:
        heap[i], heap[largest] = heap[largest], heap[i]
        max_heapify(heap, heap_size, largest)


def min_heapify(heap, heap_size, i):
    l = left(i)
    r = right(i)

    if l <= heap_size and heap[l] < heap[i]:
        smallest = l
    else:
        smallest = i

    if r <= heap_size and heap[r] < heap[smallest]:
        smallest = r

    if i != smallest:
        heap[i], heap[smallest] = heap[smallest], heap[i]
        min_heapify(heap, heap_size, smallest)


def insert_node_max_heap(heap, heap_size, node):
    heap_size += 1
    heap[heap_size] = node
    i = heap_size
    while i > 1 and heap[i] > heap[parent(i)]:
        heap[parent(i)], heap[i] = heap[i], heap[parent(i)]
        i = parent(i)
    return heap_size


def insert_node_min_heap(heap, heap_size, node):
    heap_size += 1
    heap[heap_size] = node
    i = heap_size
    while i > 1 and heap[i] < heap[parent(i)]:
        heap[parent(i)], heap[i] = heap[i], heap[parent(i)]
        i = parent(i)

    return heap_size


def extract_max(heap, heap_size):
    max_item = heap[1]
    heap[1] = heap[heap_size]
    heap_size -= 1
    max_heapify(heap, heap_size, 1)
    return max_item, heap_size


def extract_min(heap, heap_size):
    min_item = heap[1]
    heap[1] = heap[heap_size]
    heap_size -= 1
    min_heapify(heap, heap_size, 1)
    return min_item, heap_size


if __name__ == "__main__":   
    min_heap_size, max_heap_size = 0, 0

    n = int(input().strip())
    max_heap = [-1] * (n + 1)
    min_heap = [-1] * (n + 1)

    a = []
    a_i = 0
    for a_i in range(n):
        a_t = int(input().strip())
        a.append(a_t)

    median = a[0]
    print("%.1f" % median)

    if len(a) < 1:
        sys.exit(0)

    second_number = a[1]

    if second_number > median:
        min_heap[1] = second_number
        min_heap_size += 1
        max_heap[1] = median
        max_heap_size += 1
    else:
        min_heap[1] = median
        min_heap_size += 1
        max_heap[1] = second_number
        max_heap_size += 1

    median = (a[0] + a[1]) / 2.0
    print("%.1f" % median)

    for num in a[1:]:
        
        if num > median:
            min_heap_size = insert_node_min_heap(min_heap, min_heap_size, num)
            if min_heap_size - max_heap_size > 1:
                node, min_heap_size = extract_min(min_heap, min_heap_size)
                max_heap_size = insert_node_max_heap(max_heap, max_heap_size, node)

        else:
            max_heap_size = insert_node_max_heap(max_heap, max_heap_size, num)
            if max_heap_size - min_heap_size > 1:
                node, max_heap_size = extract_max(max_heap, max_heap_size)
                min_heap_size = insert_node_min_heap(min_heap, min_heap_size, node)

        #print(max_heap)
        #print(min_heap)

        if min_heap_size > max_heap_size:
            median = min_heap[1]
        elif max_heap_size > min_heap_size:
            median = max_heap[1]
        else:
            median = (min_heap[1] + max_heap[1]) / 2.0
            
        print("%.1f" % median)