A data structure designed to perform median finding in constant time. This implementation is part of the solution for Question 1, Homework 4, in the CSCI 570 course during the Spring 2024 semester at the University of Southern California.

median_heap.cpp

/*
 * C++20 Standard Required
 * Author: Sicheng He
 * Date: 2024-02-06
 */
#include<iostream>
#include<queue>
#include<vector>
#include<concepts>
#include<optional>
#include<format>

template <typename T>
concept NumericType = std::integral<T> || std::floating_point<T>;

template<NumericType T>
class MedianHeap {
private:
	std::priority_queue<T> max_heap;
	std::priority_queue<T, std::vector<T>, std::greater<T>> min_heap;
public:
	MedianHeap() = default;
	size_t size() const { return max_heap.size() + min_heap.size(); }
	std::optional<T> find_median() const {
		if (size() == 0) {
			return std::nullopt;
		}
		if (max_heap.size() == min_heap.size()) {
			return (max_heap.top() + min_heap.top()) / 2;
		}
		else if (max_heap.size() < min_heap.size()) {
			return min_heap.top();
		}
		else {
			return max_heap.top();
		}
	}
	void insert(T element) {
		if (max_heap.empty()) {
			max_heap.push(element);
			return;
		}
		if (min_heap.empty()) {
			min_heap.push(element);
			return;
		}

		if (element < max_heap.top()) {
			max_heap.push(element);
			if (max_heap.size() >= min_heap.size() + 2) {
				int tmp = max_heap.top();
				max_heap.pop();
				min_heap.push(tmp);
			}
		}
		else if (element > min_heap.top()) {
			min_heap.push(element);
			if (min_heap.size() >= max_heap.size() + 2) {
				int tmp = min_heap.top();
				min_heap.pop();
				max_heap.push(tmp);
			}
		}
		else {
			if (max_heap.size() < min_heap.size()) {
				max_heap.push(element);
			}
			else {
				min_heap.push(element);
			}
		}
	}

};

int main() {
	auto median_heap = MedianHeap<float>{};
	// []
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(1);
	// [1]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(666);
	// [1, 666]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(124);
	// [1, 124, 666]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(12356);
	// [1, 124, 666, 12356]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(12356);
	// [1, 124, 666, 12356, 12356]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(2);
	// [1, 2, 124, 666, 12356, 12356]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;
	median_heap.insert(789);
	// [1, 2, 124, 666, 789, 12356, 12356]
	std::cout << std::vformat("Median: {}", std::make_format_args(median_heap.find_median().value_or(-114514))) << std::endl;


}