20-Coding-Patterns-DesignGurus.org/10_find_the_median_of_a_number_stream.cpp

/* Problem: https://leetcode.com/problems/find-median-from-data-stream/
 */

#include <algorithm>
#include <functional>
#include <iostream>
#include <limits>
#include <queue>
#include <vector>

enum class action { addNum, findMedian };

struct instruction {
  action do_;
  double val;
};
class MedianFinder {
private:
  std::priority_queue<double, std::vector<double>, std::greater<double>>
      min_heap{std::greater<double>{}, {std::numeric_limits<double>::max()}};
  std::priority_queue<double> max_heap{std::less<double>{},
                                       {std::numeric_limits<double>::lowest()}};

public:
  MedianFinder() {}

  void addNum(double num) {
    if (min_heap.size() == max_heap.size()) {
      if (num > min_heap.top()) {
        min_heap.push(num);
      } else {
        max_heap.push(num);
      }
    } else if (min_heap.size() < max_heap.size()) {
      if (num > max_heap.top()) {
        min_heap.push(num);
      } else {
        min_heap.push(max_heap.top());
        max_heap.pop();
        max_heap.push(num);
      }
    } else {
      if (num > min_heap.top()) {
        max_heap.push(min_heap.top());
        min_heap.pop();
        min_heap.push(num);
      } else {
        max_heap.push(num);
      }
    }
  }

  double findMedian() {
    if (min_heap.size() == max_heap.size()) {
      return (min_heap.top() + max_heap.top()) / 2;
    } else if (min_heap.size() < max_heap.size()) {
      return max_heap.top();
    } else {
      return min_heap.top();
    }
  }
};

int main() {
  std::vector<instruction> instructions = {{action::addNum, 1},
                                           {action::addNum, 2},
                                           {action::findMedian, 0},
                                           {action::addNum, 3},
                                           {action::findMedian, 0}};
  MedianFinder *median_finder = new MedianFinder();
  std::cout << "{";
  for (instruction &i : instructions) {
    if (i.do_ == action::addNum) {
      median_finder->addNum(i.val);
      std::cout << "null, ";
    } else if (i.do_ == action::findMedian) {
      i.val = median_finder->findMedian();
      std::cout << i.val << ", ";
    }
  }
  std::cout << "\b\b}\n";
}