Efficient Median Calculation for Array Prefixes Using Heaps | CodeSignal Learn

Efficient Median Calculation for Array Prefixes Using Heaps | CodeSignal Learn

Introduction

Hello there, budding programmer! I hope you're ready, because today, we're going to dive deep into high-level data manipulation and increase our understanding of heaps . Heaps are fundamental data structures commonly used in algorithms. We're going to leverage their potential today in an interesting algorithmic problem. Are you ready for the challenge? Let's get started!

Task Statement

We have a task at hand related to array manipulation and the use of heaps. The task is as follows: Given a list of unique integers with elements ranging from 1 to 10^6 and length between 1 to 1000, we need to create a Python function prefix_median() . This function will take the list as an input and return a corresponding list, which consists of the medians of all the prefixes of the input list. Remember that a prefix of an array is a contiguous subsequence that starts from the first element. And the median of a sequence of numbers is the middle number when the sequence is sorted. If the length of the sequence is even, the median is half the sum of the middle two elements. For example, consider an input list [1, 9, 2, 8, 3] . The output of your function should be [1, 5.0, 2, 5.0, 3] .

Heap and Its Operations

A Heap is a useful tool in Python that helps efficiently organize and retrieve data based on their values. In our context, we use a specific type of heap called a Min Heap, where the smallest element is located at the beginning of our heap structure. For our task, we'll be using these principal operations: Adding Elements : You can add a new element to a heap using heapq.heappush(heap, item) . This operation ensures the element is correctly positioned to maintain the Min Heap property. Removing Elements : heapq.heappop(heap) is used to remove and return the smallest element from the heap. The heap readjusts itself to sustain the Min Heap property. Accessing Minimum Element : If you want to inspect the smallest element without removing it from the heap, you can simply refer to the first position of the heap as heap[0] . These operations ensure that the smallest element can always be gathered quickly, at constant time O(1), and new elements can be added maintaining the heap structure at logarithmic time, O(log n).

Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Alright, let's break our approach down into manageable steps. To begin with, we're going to need two heaps: a min heap to store the larger half of the numbers seen so far, and a max heap to store the smaller half. We'll also need a list to store the median for each prefix. As you know, a heap is a binary tree in which a parent node is some sort of extreme (either minimum or maximum) of its children. Now, let's initialize these. def prefix_median(arr): min_heap, max_heap = [], [] medians = []

Solution Building: Step 2

As the next step, we will sequentially take each number from the list and, depending on its value, push it into the min_heap or the max_heap . If it is smaller than the maximum of the lower half, it will go into the max_heap . Otherwise, it will go into the min_heap . import heapq def prefix_median(arr): min_heap, max_heap = [], [] medians = [] for num in arr: if max_heap and num < -max_heap[0]: heapq.heappush(max_heap, -num) else: heapq.heappush(min_heap, num)

Solution Building: Step 3

Next, we need to balance the two heaps to ensure that the difference between their sizes is never more than one. This way, we can always have quick access to the median. If the max_heap size becomes larger than the min_heap, we pop the max_heap's top element and push it to the min_heap . If the min_heap becomes more than one element larger than the max_heap, we do the reverse. import heapq def prefix_median(arr): min_heap, max_heap = [], [] medians = [] for num in arr: if max_heap and num < -max_heap[0]: heapq.heappush(max_heap, -num) else: heapq.heappush(min_heap, num) if len(max_heap) > len(min_heap): heapq.heappush(min_heap, -heapq.heappop(max_heap)) elif len(min_heap) > len(max_heap) + 1: heapq.heappush(max_heap, -heapq.heappop(min_heap))

Solution Building: Step 4

Having balanced the heaps, we've set ourselves up for the effortless retrieval of the median. We compute the median based on the elements at the top of the max_heap and min_heap, and then append it to our list of medians. import heapq def prefix_median(arr): min_heap, max_heap = [], [] medians = [] for num in arr: if max_heap and num < -max_heap[0]: heapq.heappush(max_heap, -num) else: heapq.heappush(min_heap, num) if len(max_heap) > len(min_heap): heapq.heappush(min_heap, -heapq.heappop(max_heap)) elif len(min_heap) > len(max_heap) + 1: heapq.heappush(max_heap, -heapq.heappop(min_heap)) if len(min_heap) == len(max_heap): medians.append((-max_heap[0] + min_heap[0]) / 2) else: medians.append(min_heap[0]) return medians Voila! You have successfully implemented a function that solves the problem using heaps in Python.

Python

import heapq


def prefix_median(arr):
    min_heap, max_heap = [], []
    medians = []
    
    for num in arr:
        if max_heap and num < -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            heapq.heappush(min_heap, num)

Python

import heapq


def prefix_median(arr):
    min_heap, max_heap = [], []
    medians = []
    
    for num in arr:
        if max_heap and num < -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            heapq.heappush(min_heap, num)
        
        if len(max_heap) > len(min_heap):
            heapq.heappush(min_heap, -heapq.heappop(max_heap))
        elif len(min_heap) > len(max_heap) + 1:
            heapq.heappush(max_heap, -heapq.heappop(min_heap))

Python

import heapq


def prefix_median(arr):
    min_heap, max_heap = [], []
    medians = []
    
    for num in arr:
        if max_heap and num < -max_heap[0]:
            heapq.heappush(max_heap, -num)
        else:
            heapq.heappush(min_heap, num)
        
        if len(max_heap) > len(min_heap):
            heapq.heappush(min_heap, -heapq.heappop(max_heap))
        elif len(min_heap) > len(max_heap) + 1:
            heapq.heappush(max_heap, -heapq.heappop(min_heap))
        
        if len(min_heap) == len(max_heap):
            medians.append((-max_heap[0] + min_heap[0]) / 2)
        else:
            medians.append(min_heap[0])
    
    return medians