Java using two Tree Sets - O(n logk)

  • 8

    Inspired by this solution. to the problem: 295. Find Median from Data Stream

    However instead of using two priority queue's we use two Tree Sets as we want O(logk) for remove(element). Priority Queue would have been O(k) for remove(element) giving us an overall time complexity of O(nk) instead of O(nlogk).

    However there is an issue when it comes to duplicate values so to get around this we store the index into nums in our Tree Set. To break ties in our Tree Set comparator we compare the index.

    public double[] medianSlidingWindow(int[] nums, int k) {
        Comparator<Integer> comparator = (a, b) -> nums[a] != nums[b] ?[a], nums[b]) : a - b;
        TreeSet<Integer> left = new TreeSet<>(comparator.reversed());
        TreeSet<Integer> right = new TreeSet<>(comparator);
        Supplier<Double> median = (k % 2 == 0) ?
            () -> ((double) nums[left.first()] + nums[right.first()]) / 2 :
            () -> (double) nums[right.first()];
        // balance lefts size and rights size (if not equal then right will be larger by one)
        Runnable balance = () -> { while (left.size() > right.size()) right.add(left.pollFirst()); };
        double[] result = new double[nums.length - k + 1];
        for (int i = 0; i < k; i++) left.add(i);; result[0] = median.get();
        for (int i = k, r = 1; i < nums.length; i++, r++) {
            // remove tail of window from either left or right
            if(!left.remove(i - k)) right.remove(i - k);
            // add next num, this will always increase left size
            right.add(i); left.add(right.pollFirst());
            // rebalance left and right, then get median from them
  ; result[r] = median.get();
        return result;

  • 0

    I have learned a lot from this post. Thank you very much.

  • 1

    Is remove(element) really O(1) for TreeSet?

    The Java API for TreeSet says,
    This implementation provides guaranteed log(n) time cost for the basic operations (add, remove and contains).

  • 1


    My bad! Good catch, it should be O(logk) for remove (will update the post). It doesn't change the overall runtime complexity. If we would have used PriorityQueue then remove(object) would be linear time instead of logarithmic, affecting the overall runtime complexity.

  • 0
    This post is deleted!

  • 0

    Definitely awesome solution. Thanks!

  • 0

    java lambdas are magical

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.