Using Heaps in C++ to Calculate Prefix Medians

Using Heaps in C++ to Calculate Prefix Medians | 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

1

Heap and Its Operations

O(1)

Join the 1M+ learners on CodeSignal

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

A heap is a useful tool in C++ that helps efficiently organize and retrieve data based on their values.

In our context, we use two specific types of heaps: a Min Heap and a Max Heap. The Min Heap is used to store the larger half of the numbers seen so far, while the Max Heap stores the smaller half. In C++, heaps are implemented as priority_queue containers.

For our task, we'll be using these principal operations:

Adding Elements: You can add a new element to a heap using the push() method. By default, C++'s priority_queue is a Max Heap. To create a Min Heap, you need to use a custom comparator like greater<int>. This ensures that the smallest element is always at the top.
C++
```
std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap;
```
Removing Elements: The pop() method removes the top element from the heap, which is the smallest element in the Min Heap or the largest element in the Max Heap.
Accessing Minimum Element: The top() method is used to inspect the root element of the heap without removing it. For the Min Heap with a greater<int> comparator, top() will return the smallest element.

These operations ensure that the root element can always be gathered quickly, in constant time $O(1)$ , and new elements can be added while maintaining the heap structure in logarithmic time, $O(log n)$ .

Solution Building: Step 1

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 vector to store the median for each prefix. Now, let's initialize these. C++#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; }#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; }

Solution Building: Step 2

As the next step, we will sequentially take each number from the vector and, depending on its value, push it into the minHeap or the maxHeap. If it is smaller than the maximum of the lower half, it will go into the maxHeap. Otherwise, it will go into the minHeap. C++#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(num); } } }#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(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 maxHeap size becomes larger than the minHeap, we pop the maxHeap's top element and push it to the minHeap. If the minHeap becomes more than one element larger than the maxHeap, we do the reverse. C++#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(num); } if (maxHeap.size() > minHeap.size()) { minHeap.push(maxHeap.top()); maxHeap.pop(); } else if (minHeap.size() > maxHeap.size() + 1) { maxHeap.push(minHeap.top()); minHeap.pop(); } } }#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(num); } if (maxHeap.size() > minHeap.size()) { minHeap.push(maxHeap.top()); maxHeap.pop(); } else if (minHeap.size() > maxHeap.size() + 1) { maxHeap.push(minHeap.top()); minHeap.pop(); } } }

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 maxHeap and minHeap, and then append it to our vector of medians. C++#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(num); } if (maxHeap.size() > minHeap.size()) { minHeap.push(maxHeap.top()); maxHeap.pop(); } else if (minHeap.size() > maxHeap.size() + 1) { maxHeap.push(minHeap.top()); minHeap.pop(); } if (minHeap.size() == maxHeap.size()) { medians.push_back(maxHeap.top()); } else { medians.push_back(minHeap.top()); } } return medians; }#include <iostream> #include <vector> #include <queue> std::vector<int> prefixMedian(const std::vector<int>& arr) { std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap; std::priority_queue<int> maxHeap; std::vector<int> medians; for (int num : arr) { if (!maxHeap.empty() && num < maxHeap.top()) { maxHeap.push(num); } else { minHeap.push(num); } if (maxHeap.size() > minHeap.size()) { minHeap.push(maxHeap.top()); maxHeap.pop(); } else if (minHeap.size() > maxHeap.size() + 1) { maxHeap.push(minHeap.top()); minHeap.pop(); } if (minHeap.size() == maxHeap.size()) { medians.push_back(maxHeap.top()); } else { medians.push_back(minHeap.top()); } } return medians; }