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