Problem: 169, Find the median of the stream

Author: mandliya

README.md

@@ -6,7 +6,7 @@
 
 | Current Status|     Stats     |
 | :------------: | :----------: |
-| Total Problems | 168 |
+| Total Problems | 169 |
 
 </center>
 
@@ -240,4 +240,5 @@ Include contains single header implementation of data structures and some algori
 |Product of Array Except Self. Given an array of n integers where n > 1, nums, return an array output such that output[i] is equal to the product of all the elements of nums except nums[i].| [product_except_self.cpp](leet_code_problems/product_except_self.cpp)|
 |Given a sorted array, remove duplicates in place and return the new length. It doesn't matter what is in array beyond the unique elements size. Expected O(1) space and O(n) time complexity.| [remove_duplicates.cpp](leet_code_problems/remove_duplicates.cpp) |
 | Count the number of islands in a grid. Given a grid representing 1 as land body, and 0 as water body, determine the number of islands (more details in problem comments)|[count_islands.cpp](leet_code_problems/count_islands.cpp)|
+| Find median from a data stream. Design a data structure that supports addNum to add a number to the stream, and findMedian to return the median of the current numbers seen so far. Also, if the count of numbers is even, return average of two middle elements, return median otherwise.|[median_stream.cpp](leet_code_problems/median_stream.cpp)
 

leet_code_problems/median_stream.cpp

@@ -0,0 +1,140 @@
+/*
+ * Source : Leetcode, 295. Find Median from Data Stream
+ * Category : hard
+ * 
+ * Find median from a data stream. Design a data structure that supports addNum to add a number to the stream,
+ * and findMedian to return the median of the current numbers seen so far.  Also, if the count of numbers is
+ * even, return average of two middle elements, return median otherwise.
+ * 
+ * In naive approach, we could maintain an internal vector and sort the array everytime median is requested,
+ * and return the middle element(when odd count) or mean of two middle elements (when mean count)
+ * 
+ * However, this doesn't scale well as the size of the vector grows, and if 'findMedian' is accessed too
+ * frequently. We would want a potentially better solution which could reduce the cost of finding median,
+ * even though we might make add a little slow.
+ * 
+ * We could use something which maintains a ordered stream of numbers, and median could be retrieved in O(1)
+ * time, and add could be done in O(logn) time.
+ * 
+ * This could be achieved using a Self-Balanced Binary search tree, or using two heaps.
+ * We will currently solve with a heap approach.
+ * 
+ * Why two heaps?
+ * We want two middle numbers (in even count case), so if we maintain two heaps, one min-heap and other
+ * max-heap, the max-heap will maintain the lower ordered half of the numbers while the min-heap would
+ * maintained the upper half of the numbers.
+ * Consider even number of numbers in stream so far, in that case, the top of the lower heap (max heap)
+ * will contain the left middle number of the stream and top of upper heap (min heap) will contain the 
+ * right middle number of the stream. Thus we could retrive the two middle numbers in O(1) time.
+ * 
+ * [1, 2, 3, 4, 5]  <-- values in lower heap. (MAX heap)
+ * [6, 7, 8, 9, 10] <-- Values in upper heap. (MIN heap)
+ * 
+ * lower.top() = 5
+ * upper.top() = 6
+ * So we will want (6+5)/2 = 5.5 as answer.
+ * 
+ * Now, the algorithm:
+ * In order to maintain the two middle numbers of the stream on the top of the two heaps,
+ * we will have to balance them with each addition of new number.
+ * 
+ * In case of odd numbers, the lower heap will contain 1 number extra then the upper heap.
+ * In case of even numbers, both the heaps will have equal numbers of elements.
+ * 
+ * As we add a new number, we will add it to lower, then we will balance the heaps by transferring
+ * the top to upper heap, and then we will maintain the size property maintained above.
+ */
+
+#include <iostream>
+#include <queue>
+#include <vector>
+#include <random>
+
+class MedianFinder {
+public:
+    MedianFinder() {}
+    void addNum(int num)
+    {
+        // first add it to lower.
+        lower.push(num);
+
+        // balance the heaps to set the right values at the top.
+        upper.push(lower.top());
+        lower.pop();
+
+        // Maintain the size property.
+        if (lower.size() < upper.size()) {
+            lower.push(upper.top());
+            upper.pop();
+        }
+    }
+
+    double findMedian()
+    {
+        return (lower.size() > upper.size()) ? static_cast<double>(lower.top()) :
+            (static_cast<double>(lower.top() + upper.top()) / 2);
+    }
+
+    // utility function to print current sorted stream.
+    void printSortedStream()
+    {
+        std::priority_queue<int> lc{lower};
+        std::priority_queue<int, std::vector<int>, std::greater<int>> uc{upper};
+        std::vector<int> temp;
+        // lower half first;
+        while(!lc.empty()){
+            temp.push_back(lc.top());
+            lc.pop();
+        }
+
+        std::reverse(temp.begin(), temp.end());
+
+        for (auto n: temp) {
+            std::cout << n << " ";
+        }
+
+        while (!uc.empty()) {
+            std::cout << uc.top() << " ";
+            uc.pop();
+        }
+
+        std::cout << std::endl;
+    }
+private:
+    // max heap maintaining the lower half of ordered stream.
+    std::priority_queue<int> lower; 
+    // min heap maintaining the upper half of ordered stream.
+    std::priority_queue<int, std::vector<int>, std::greater<int>> upper;
+};
+
+
+int main()
+{
+    std::default_random_engine generator;
+    std::uniform_int_distribution<int> distribution(0,20);
+    MedianFinder medianFinder;
+
+    for (int i=0; i<5; ++i)
+    {
+        medianFinder.addNum(distribution(generator));
+    }
+
+    // current state of stream
+    medianFinder.printSortedStream();
+    std::cout << "Median of current stream: " << medianFinder.findMedian()
+        << std::endl;
+
+    // Add 5 more.
+    //
+    for (int i=0; i<5; ++i)
+    {
+        medianFinder.addNum(distribution(generator));
+    }
+    
+     // current state of stream
+    medianFinder.printSortedStream();
+    std::cout << "Median of current stream: " << medianFinder.findMedian()
+        << std::endl;
+
+    return 0;
+}