Very Short, O(log n) + O(1)


  • 32

    Same idea as before, but really exploiting the symmetry of the two heaps by switching them whenever a number is added. Still O(log n) for adding and O(1) for median. Partially inspired by peisi's updated solution.

    Update: Added a new Java version (the first one).


    Java

    class MedianFinder {
        
        Queue<Integer> q = new PriorityQueue(), z = q, t,
                       Q = new PriorityQueue(Collections.reverseOrder()); 
    
        public void addNum(int num) {
            (t=Q).add(num);
            (Q=q).add((q=t).poll());
        }
    
        public double findMedian() {
            return (Q.peek() + z.peek()) / 2.;
        }
    };
    

    Or:

    class MedianFinder {
    
        Queue[] q = {new PriorityQueue(), new PriorityQueue(Collections.reverseOrder())};
        int i = 0;
    
        public void addNum(int num) {
            q[i].add(num);
            q[i^=1].add(q[i^1].poll());
        }
    
        public double findMedian() {
            return ((int)(q[1].peek()) + (int)(q[i].peek())) / 2.0;
        }
    };
    

    Python

    from heapq import *
    
    class MedianFinder:
    
        def __init__(self):
            self.heaps = None, [], []
            self.i = 1
    
        def addNum(self, num):
            heappush(self.heaps[-self.i], -heappushpop(self.heaps[self.i], num * self.i))
            self.i *= -1
    
        def findMedian(self):
            return (self.heaps[self.i][0] * self.i - self.heaps[-1][0]) / 2.0
    

    Or:

    from heapq import *
    
    class MedianFinder:
    
        def __init__(self):
            self.data = 1, [], []
    
        def addNum(self, num):
            sign, h1, h2 = self.data
            heappush(h2, -heappushpop(h1, num * sign))
            self.data = -sign, h2, h1
    
        def findMedian(self):
            sign, h1, h2 = d = self.data
            return (h1[0] * sign - d[-sign][0]) / 2.0

  • 1
    L

    Shocked again by your using of q[i^=1].add(q[i^1].poll()); and return ((int)(q[1].peek()) + (int)(q[i].peek())) / 2.0;. It's art!


  • 0

    If only Java supported generic arrays so I wouldn't need those (int)(...) casts. Sigh...


  • 0
    X

    awesome!

    I was thinking maybe we can use AVL tree, and the node of tree should has attributes: # of nodes in left child and # of nodes in right child, the time complexity should also be O(log n) + O(1), but it's would be very hard to implement AVL tree right in interview.


  • 0

    Wouldn't findMedian take O(log n) as well? AVL trees are only roughly balanced, so the one or two middle values could be at the bottom of the tree, no?


  • 0
    X

    ah, right. I was wrong


  • 0
    This post is deleted!

  • 0

    Have a look at my new first Java solution I just added, maybe I can shock you once more :-)


  • 0
    M

    Yep, I'm shocked.


  • 0
    A

    Your solution looks awesome, but I am not able to understand the syntax.
    Help Help.


  • 0
    D

    oooo... another code to debug to understand how it works :)
    I didn't know that this syntax returns tuple.

    >>> heaps = None, [], []
    >>> heaps
    (None, [], [])
    

  • 1
    D

    I'm not sure what could be the reaction of interviewer if I code this way. Anyway I like shorter more. :) Thx


  • 0
    W

    UNBELIEVABLY BEAUTIFUL SOLUTION!


  • 0

    If you can write code in C# as short as this one then you are GOAT


  • 0
    W

    @StefanPochmann There are only min-heap and max-heap in your original solution but there are four here. Could you elaborate more on how the variables z and t are used?
    I do not quite understand the assignments in (t=Q).add(num); and (Q=q).add((q=t).poll());.
    Also, why z.peek() is used and it works for sizes of both odd and even?


  • 0

    @wai_ting said in Very Short, O(log n) + O(1):

    @StefanPochmann There are only min-heap and max-heap in your original solution but there are four here.

    Huh? No, there are also only two here.

    Could you elaborate more on how the variables z and t are used?
    I do not quite understand the assignments in (t=Q).add(num); and (Q=q).add((q=t).poll());.
    Also, why z.peek() is used and it works for sizes of both odd and even?

    t is temporary, just used for swapping q and Q with the usual three-way assignment. The only unusual thing is that I embedded the assignments into the other code.

    I think I meant z as "zero", as in the index of the solution using Queue[] q. The first queue I add into. So I can still access that one regardless of how often I swapped q and Q.


  • 1
    W

    @StefanPochmann Thanks a lot for your explanation. It helps me understand your solution better. I tried to run your solution steps by steps and it looks brillant. Basically q and Q point to min-heap and max-heap alternatively and z points to Q for odd size and to q for even size.

    For example, add and find median for [3,4,1,5,2]

    q Q t z Median Remark
    Initially [] [] null [] 0 q: min-heap, Q: max-heap
    After 1st iteration [] [3] q Q 3.0 q: max-heap, Q: min-heap
    After 2nd iteration [4] [3] q q 3.5 q: min-heap, Q: max-heap
    After 3rd iteration [1] [3,4] q Q 3.0 q: max-heap, Q: min-heap
    After 4th iteration [4,5] [3,1] q q 3.5 q: min-heap, Q: max-heap
    After 5th iteration [2,1] [3,5,4] q Q 3.0 q: max-heap, Q: min-heap

    Expanding addNum would look like this:

    public void addNum(int num) {
      t = Q;
      t.add(num);
      Q = q;
      q = t;
      Q.add(q.poll());
    }
    

  • 0

    @wai_ting If you "expand" it, I'd do it like this, cleanly do the swap after the actual work and use t only there.

    public void addNum(int num) {
        Q.add(num);
        q.add(Q.poll());
        t = Q;
        Q = q;
        q = t;
    }

  • 0
    W

    @StefanPochmann Thanks for your suggestion. It definitely looks cleaner than mine.

    Just curious, I tried to swap Q and q by a call swap(Q, q); with the following implementation instead of inlining the swap. I got a null pointer exception when finding the median. I suppose it has something to do with what z is pointing to. Any idea what the differences are?

    private void swap(Queue<Integer> q1, Queue<Integer> q2) {
        t = q1;
        q1 = q2;
        q2 = q1;
    }
    

  • 0

    @wai_ting Your swap function doesn't work. You're just swapping the swap function's own variables. It has no effect for the caller.


Log in to reply
 

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