【LeetCode】295. 数据流的中位数（同剑指Offer41）

一、题目

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

• void addNum(int num) - 从数据流中添加一个整数到数据结构中。
• double findMedian() - 返回目前所有元素的中位数。
  示例：

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

进阶:

1. 如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
2. 如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？

二、解决

1、暴力

思路：

比较简单，直接看代码即可。比较慢，不推荐。

代码：

class MedianFinder {

    List<Integer> res = new ArrayList<>();
    
    public void addNum(int num) {
        res.add(num);
    }
    public double findMedian() {
        int n = res.size();
        Collections.sort(res);
        return n%2==1 ? res.get(n/2) : (res.get(n/2-1)+res.get(n/2))/2.0;
    }
}

时间复杂度： $O(nlogn)$
空间复杂度： $O(n)$

2、插入

思路：

将数加入一个有序数组，然后返回中间位置。具体请结合注释理解。比较慢，不推荐。

代码：

class MedianFinder {

    List<Integer> res = new ArrayList<>();
    
    public void addNum(int num) {
    	// 二分查找 O(logn)
    	...
    	// 挪动位置，然后插入 O(n)
    	...
        res.add(num);
    }
    public double findMedian() {
        int n = res.size();
        return n%2==1 ? res.get(n/2) : (res.get(n/2-1)+res.get(n/2))/2.0;
    }
}

时间复杂度： $O(n)$
空间复杂度： $O(n)$

3、大小堆

思路： （主要来自参考2）

1、中位数位置特点

对数据流排序后，中位数位于这些数中间1个（奇数个）或两个数取平均（偶数个）。

2、中位数值特点

中位数只位于前有序数组的最大值（奇数个），或者是 （前有序数组的最大值+后有序数组的最小值）/ 2。

即我们只关心两个有序数组的最值，所以可以用两个优先队列来解决。前面只关心最大值，用大顶堆，后面只关心最小值，用小顶堆。

3、两个堆性质

1）大顶堆堆顶元素 <= 小顶堆元素；
2）大顶堆元素个数与小顶堆元素相等，或者多1.

说明：maxHeap.size() +1 == minHeap.size()  X
      即大顶堆元素个数比小顶堆元素个数少1的情况不可能出现，因为后有序数组元素个数不会多余前有序数组元素个数

4、数量控制

1）maxHeap.size() == minHeap.size() +1

• 为让大顶堆元素多1个，可采用流程：大顶堆 --> 小顶堆 --> 大顶堆

2）maxHeap.size() == minHeap.size()

• 在前一步，大顶堆元素多一个，现在添加1个元素，为了使得两者相等，必须要有一个元素移入小顶堆，即必须具有流程：大顶堆 --> 小顶堆。

小结： 一定要有 大顶堆–>小顶堆 的流程，为了保证大顶堆的元素个数 >= 小顶堆的元素个数，后面再根据情况决定是否加上 小顶堆–>大顶堆 的流程。

代码：

PriorityQueue<Integer> minHeap = new PriorityQueue<>();//heap is a minimal heap by default
PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());//change to a maximum heap
// Other maxHeap define method
// PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> (b - a));

// Adds a number into the data structure.
public void addNum(int num) {
	maxHeap.offer(num);
	minHeap.offer(maxHeap.poll());
	if (maxHeap.size() < minHeap.size())
		maxHeap.offer(minHeap.poll());
}

// Returns the median of current data stream
public double findMedian() {
	if (maxHeap.size() == minHeap.size())
		return (maxHeap.peek() + minHeap.peek()) / 2.0;
	else
		return maxHeap.peek();
}

时间复杂度： $O(logn)$，插入$O(logn)$，查询$O(1)$
空间复杂度： $O(n)$

4、BST

思路：

用二分搜索树来解决上述问题。这里的TreeNode声明加了一个size字段。其他的主要还是递归的方法，理解不算难。

代码-递归版：

class MedianFinder {

	private class TreeNode {
		int val;
		int size;
		TreeNode left;
		TreeNode right;

		TreeNode(int val) {
			this.val = val;
			this.size = 1;
		}
	}

	private TreeNode root;

	/** initialize your data structure here. */
	public MedianFinder() {

	}

	public void addNum(int num) {
		if (this.root == null) {
			root = new TreeNode(num);
		} else {
			this.addNum(num, this.root);
		}
	}

	private void addNum(int num, TreeNode node) {
		if (node.val >= num) {
			if (node.left == null) {
				node.left = new TreeNode(num);
			} else {
				addNum(num, node.left);
			}
		} else {
			if (node.right == null) {
				node.right = new TreeNode(num);
			} else {
				addNum(num, node.right);
			}
		}
		node.size++;
	}

	public double findMedian() {
		if (this.root.size % 2 == 1) {
			return search(this.root.size / 2 + 1, this.root);
		} else {
			int left = search(this.root.size / 2, this.root);
			int right = search(this.root.size / 2 + 1, this.root);
			return (left + right) / 2.0;
		}
	}

	public int search(int index, TreeNode node) {
		if (node.size == 1) {
			return node.val;
		}
		int leftSize = node.left != null ? node.left.size : 0;
		if (leftSize >= index) {
			return search(index, node.left);
		} else if (leftSize + 1 == index) {
			return node.val;
		} else {
			return search(index - leftSize - 1, node.right);
		}
	}
}

思路： 类似上面，只不过用了迭代的方式。

代码-迭代版：

class MedianFinder {
    class TreeNode{
        int val;
        TreeNode parent,left,right;
        TreeNode(int val, TreeNode p){
            this.val=val;
            this.parent=p;
            left=null;
            right=null;
        }
        void add(int num){
            if(num>=val){
                if(right==null)
                    right=new TreeNode(num,this);
                else
                    right.add(num);
            }else{
                if(left==null)
                    left=new TreeNode(num,this);
                else
                    left.add(num);
            }
        }
        TreeNode next(){
            TreeNode ret;
            if(right!=null){
                ret=right;
                while(ret.left!=null)
                    ret=ret.left;
            }else{
                ret=this;
                while(ret.parent.right==ret)
                    ret=ret.parent;
                ret=ret.parent;
            }
            return ret;
        }
        TreeNode prev(){
            TreeNode ret;
            if(left!=null){
                ret=left;
                while(ret.right!=null)
                    ret=ret.right;
            }else{
                ret=this;
                while(ret.parent.left==ret)
                    ret=ret.parent;
                ret=ret.parent;
            }
            return ret;
        }
    }
    int n;
    TreeNode root, curr;
    // Adds a number into the data structure.
    public void addNum(int num) {
        if(root==null){
            root = new TreeNode(num,null);
            curr=root;
            n=1;
        }else{
            root.add(num);
            n++;
            if (n%2==1) {
                if(curr.val<=num)
                curr = curr.next();
            } else if(curr.val>num)
                curr = curr.prev();
        }
    }

    // Returns the median of current data stream
    public double findMedian() {
        if(n%2==0){
            return ((double)curr.next().val+curr.val)/2;
        }else
            return curr.val;
    }
};

时间复杂度： $O(logn)$
空间复杂度： $O(n)$

5、提问

// 不能完全通过
class MedianFinder {
    int A[] = new int[101], n = 0;
    
    public void addNum(int num) {
        A[num]++;
        n++;
    }
    
    public double findMedian() {
        int count = 0, i = 0;
        while (count < n/2) count += A[i++];
        int j = i;
        while (count < n/2+1) count += A[j++];
        return (n%2 == 1) ? j-1 : (i+j-2) / 2.0;
    }
}

1. If all integer numbers from the stream are between 0 and 100, how would you optimize it?

We can maintain an integer array of length 100 to store the count of each number along with a total count. Then, we can iterate over the array to find the middle value to get our median.

Time and space complexity would be O(100) = O(1).

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

In this case, we need an integer array of length 100 and a hashmap for these numbers that are not in [0,100].

三、参考

1、数据流的中位数
2、优先队列（Python 代码、Java 代码）
3、Short simple Java/C++/Python, O(log n) + O(1)
4、JAVA-----------Easy Version To Understand!!!
5、18ms beats 100% Java Solution with BST
6、22ms Java solution using binary tree, beats 99.82% of submissions
7、啥也不说了咱们来一个手撕索引重复红黑树吧。。。。。。。。。我哭了
8、Solutions to follow-ups
9、[Java] Simple Code - Follow Up