// uppi.cc -*- mode:c++, compile-command:"g++ -Wall -Wextra -pedantic -std=c++0x uppi.cc -o uppi && ./uppi" -*-
#include <iostream>
#include <algorithm>
#include <vector>
#include <set>
#include <cassert>
#include <iterator>
template <typename T>
class median_set {
public:
std::multiset<T> below, above;
// O(log N)
void rebalance()
{
int diff = above.size() - below.size();
if (diff > 0) {
below.insert(*above.begin());
above.erase(above.begin());
} else if (diff < -1) {
above.insert(*below.rbegin());
below.erase(below.find(*below.rbegin()));
}
}
public:
// O(1)
bool empty() const { return below.empty() && above.empty(); }
// O(1)
T const& median() const
{
assert(!empty());
return *below.rbegin();
}
// O(log N)
void insert(T const& value)
{
if (!empty() && value > median())
above.insert(value);
else
below.insert(value);
rebalance();
}
// O(log N)
void erase(T const& value)
{
if (value > median())
above.erase(above.find(value));
else
below.erase(below.find(value));
rebalance();
}
};
int main()
{
median_set<int> ms;
std::vector<int> v;
srand(0);
for (int i=0; i<100000; ++i) {
bool remove = !(rand()%3);
if (remove && !v.empty()) {
unsigned rnd = rand();
int idx = rnd%v.size();
ms.erase(v[idx]);
v.erase(v.begin()+idx);
} else {
int value = rand();
ms.insert(value);
v.push_back(value);
}
if (!v.empty()) {
std::sort(v.begin(), v.end());
int median1 = v[(v.size()-1)/2];
int median2 = ms.median();
assert(median1 == median2);
}
}
}