Instantly share code, notes, and snippets.

Embed
What would you like to do?
implementation of max heap and min heap to solve running median problem
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment