Java TreeSet solution with two pointers, and explanation
For some reason, there is a lack of well documented solution using TreeSet for Java developers in the discussion. Official LC solution uses C++ multiset and ite
For some reason, there is a lack of well documented solution using TreeSet for Java developers in the discussion. Official LC solution uses C++ multiset and iteration there gets easier because of pointers, which is hard to achieve using iterators in Java (considering there is no *peek* functionality in Java iterators).\n\nThe following code uses two pointer approach, and becomes a bit long because of using a custom comparator. The idea here is: instead of an iterator, we store both the number and its location in data stream in the TreeSet. This lets us handle cases where identical values are entered in the stream, and closely follows the height balanced trees approach presented in LC solution.\n\nThis is an important technique to grasp in case interviewer asks for alternative solutions to the more popular Priority Queue approach.\n\nTime complexity for add is O(log(n)) to account for the time to `add` and search (`lower` and `upper` methods of TreeSet).\nSpace complexity is O(n) since we are storing all elements of the stream here\n\n```\nclass MedianFinder {\n    TreeSet<int[]> data;\n    Comparator<int[]> customComparator;\n    int[] lower, upper;\n    int index;\n    /** initialize your data structure here. */\n    public MedianFinder() {\n        // Comparator compares first by value of the integer, then by location in stream\n        customComparator = ((int[] a, int[] b) -> {\n            if (a[0] != b[0]) {\n                return Integer.compare(a[0], b[0]);\n            }\n            else {\n                return Integer.compare(a[1], b[1]);\n            }\n        });\n        data = new TreeSet<>(customComparator);\n        lower = null;\n        upper = null;\n        index = 0;\n    }\n    \n    public void addNum(int num) {\n        int[] curr = new int[] {num, index++};\n        data.add(curr);\n        if (data.size() == 1) {\n            // First addition to the stream\n            lower = data.first();\n            upper = data.first();\n        }\n        else if (customComparator.compare(lower, upper) == 0) {\n            // last time it was odd, now it is even\n            if (customComparator.compare(curr, lower) < 0) {\n                // num is less than the left median\n                // Median has shifted one place left\n                upper = lower;\n                lower = data.lower(lower);\n            }\n            else {\n                // num is greater than or eaual to the right median\n                // Median has shifted one place right\n                lower = upper;\n                upper = data.higher(upper);\n            }\n        }\n        else {\n            // last time it was even, this time it is odd\n            // Three cases\n            if (customComparator.compare(curr, lower) < 0) {\n                // Case 1: new number inserted to the left of lower\n                // lower median is the current median in that case\n                upper = lower;\n            }\n            else  if (customComparator.compare(curr, upper) > 0) {\n                // Case 2: new number inserted to the right of upper\n                // upper median is the current median in that case\n                lower = upper;\n            }\n            else{\n                // Case 3: Either num is equal to one or both of upper, lower, or lies in between\n                // in that case, num IS our median\n                lower = curr;\n                upper = curr;\n            }\n        }\n        \n    }\n    \n    public double findMedian() {\n        if (lower == null) {\n            // Safety check\n            return 0.0;\n        }\n        // This function is trivial since we are already tracking two pointers\n        return ((double)lower[0] + upper[0])/2;\n    }\n}\n```\n\nNote: inspired by https://leetcode.com/problems/find-median-from-data-stream/discuss/439595/Java-TreeSet-Solution-beats-93, but wanted to add more commented and clean code.
['Tree', 'Ordered Set', 'Java']