295. Find Median from Data Stream

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

For example,
[2,3,4], the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

void addNum(int num) - Add a integer number from the data stream to the data structure.
double findMedian() - Return the median of all elements so far.

Example:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3)
findMedian() -> 2

Follow up:

If all integer numbers from the stream are between 0 and 100, how would you optimize it?
If 99% of all integer numbers from the stream are between 0 and 100, how would you optimize it?

//125ms
class MedianFinder {
    PriorityQueue<Integer> lowerHalf;
    PriorityQueue<Integer> upperHalf;

    /** initialize your data structure here. */
    public MedianFinder() {
        lowerHalf = new PriorityQueue<>((a,b) -> (b-a));
        upperHalf = new PriorityQueue<>();
    }

    public void addNum(int num) {
        lowerHalf.offer(num);
        upperHalf.offer(lowerHalf.poll());
        if (lowerHalf.size() < upperHalf.size()) {
            lowerHalf.offer(upperHalf.poll());
        }
    }

    public double findMedian() {
        if (upperHalf.size() == lowerHalf.size()) {
            return (upperHalf.peek() + lowerHalf.peek()) / 2.0;
        } else {
            return lowerHalf.peek() * 1.0;
        }
    }
        
}

解法2

Some notes: This solution uses an ordinary binary tree for simplicity’s sake, which means it is likely to be unbalanced. Given enough time one may well use a balanced binary tree implementation to guarantee O(logn) runtime for addNum(). It is easy to see that findMedian() runs in O(1).

By using a binary tree, we can easily keep the input numbers in nondecreasing order. Observe that whenever a number is added, the numbers used to calculate the median never shift by more than 1 position (in an imagined array representation) to the left or to the right. Let’s see an example:
[2], number used to calculate median is 2.
[2,3], numbers used are 2,3 (expanding 1 to right)
[0,2,3], use 2 (shrinking 1 to left)
[0,1,2,3], use 1,2 (expanding 1 to left)
[0,1,2,3,4], use 2 (shrinking 1 to right)
…and so on.

With this observation, in MedianFinder I employ 2 variables medianLeft and medianRight to keep track of numbers we need to calculate the median. When size is odd, they point to the same node, otherwise they always point to 2 nodes which have predecessor/successor relationship. When adding a node, we just need to check the size of our MedianFinder tree, then depending on whether the new number is inserted to the left, inbetween, or to the right of our 2 median trackers, we will change medianLeft and medianRight to point to the correct nodes. Because the position never shifts more than 1, we can simply use predecessor() or successor() on the desired node to update it. Those 2 methods run in O(logn) when the tree is balanced, hence the O(logn) runtime of addNum().

//90ms
public class MedianFinder {
    private Node root;
    private Node medianLeft;
    private Node medianRight;
    private int size;
    
    public MedianFinder() {
    }

    // Adds a number into the data structure.
    public void addNum(int num) {
        if (root == null) {
            root = new Node(num);
            medianLeft = root;
            medianRight = root;
        }
        else {
            root.addNode(num);
            if (size % 2 == 0) {
                if (num < medianLeft.data) {
                    medianRight = medianLeft;
                }
                else if (medianLeft.data <= num && num < medianRight.data) {
                    medianLeft = medianLeft.successor();
                    medianRight = medianRight.predecessor();
                }
                else if (medianRight.data <= num) {
                    medianLeft = medianRight;
                }
            }
            else {
                if (num < medianLeft.data) {
                    medianLeft = medianLeft.predecessor();
                }
                else {
                    medianRight = medianRight.successor();
                }
            }
        }
        size++;
    }

    // Returns the median of current data stream
    public double findMedian() {
        return (medianLeft.data + medianRight.data) / 2.0;
    }
    
    class Node {
        private Node parent;
        private Node left;
        private Node right;
        private int data;
        
        public Node(int data) {
            this.data = data;
        }
        
        public void addNode(int data) {
            if (data >= this.data) {
              if (right == null) {
                right = new Node(data);
                right.parent = this;
              }
              else
                right.addNode(data);
            }
            else {
              if (left == null) {
                left = new Node(data);
                left.parent = this;
              }
              else
                left.addNode(data);
            }
        }
        
        public Node predecessor() {
            if (left != null)
                return left.rightMost();
            
            Node predecessor = parent;
            Node child = this;
            
            while (predecessor != null && child != predecessor.right) {
                child = predecessor;
                predecessor = predecessor.parent;
            }
            
            return predecessor;
        }
        
        public Node successor() {
            if (right != null)
                return right.leftMost();
            
            Node successor = parent;
            Node child = this;
            
            while (successor != null && child != successor.left) {
                child = successor;
                successor = successor.parent;
            }
            
            return successor;
        }
        
        public Node leftMost(){
            if (left == null)
                return this;
            return left.leftMost();
        }
        
        private Node rightMost() {
            if (right == null)
                return this;
            return right.rightMost();
        }
        
    }
};