Median Set
(misc/median_set.hpp)

• View this file on GitHub
• Last update: 2024-01-08 00:27:17+09:00
• Include: #include "misc/median_set.hpp"

Code

#pragma once
#include <set>

/**
 * @brief Median Set
 */
template <typename T>
class MedianSet {
   public:
    void insert(T x) {
        if (x < *right.begin()) {
            left.insert(x);
        } else {
            right.insert(x);
        }
        balance();
    }

    void erase(T x) {
        if (x < *right.begin()) {
            left.erase(x);
        } else {
            right.erase(x);
        }
        balance();
    }

    T median() const {
        if (left.size() > right.size()) {
            return *left.rbegin();
        } else {
            return (*left.rbegin() + *right.begin()) / 2;
        }
    }

   private:
    std::set<T> left, right;

    void balance() {
        while (left.size() < right.size()) {
            left.insert(*right.begin());
            right.erase(right.begin());
        }
        while (left.size() > right.size() + 1) {
            right.insert(*left.rbegin());
            left.erase(left.rbegin());
        }
    }
};

#line 2 "misc/median_set.hpp"
#include <set>

/**
 * @brief Median Set
 */
template <typename T>
class MedianSet {
   public:
    void insert(T x) {
        if (x < *right.begin()) {
            left.insert(x);
        } else {
            right.insert(x);
        }
        balance();
    }

    void erase(T x) {
        if (x < *right.begin()) {
            left.erase(x);
        } else {
            right.erase(x);
        }
        balance();
    }

    T median() const {
        if (left.size() > right.size()) {
            return *left.rbegin();
        } else {
            return (*left.rbegin() + *right.begin()) / 2;
        }
    }

   private:
    std::set<T> left, right;

    void balance() {
        while (left.size() < right.size()) {
            left.insert(*right.begin());
            right.erase(right.begin());
        }
        while (left.size() > right.size() + 1) {
            right.insert(*left.rbegin());
            left.erase(left.rbegin());
        }
    }
};

Back to top page