[LeetCode困难] 295. 数据流的中位数

发表于 评论数：

记录一些做题过程中碰到的巧妙解法。

题目描述：

problem

解法一：插入排序

首先想到的办法就是：每次添加数据的时候走一个插入排序，这样我的数据结构里就维护的是一个有序数列，那么取中位数的时候直接判断一下长度的奇偶然后相应返回即可。

数据结构里面可以用ArrayList或数组保存数据。我首先实现的数组的版本，考虑到数组的可扩容性，维护一个top变量和capacity变量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class MedianFinder {

    private int[] data;
    private int top;
    private int capacity;

    /** initialize your data structure here. */
    public MedianFinder() {
        capacity = 10000;
        top = 0;
        data = new int[capacity];
    }

    public void addNum(int num) {
        if(top == capacity){
            capacity*=2;
            data = Arrays.copyOf(data, capacity);
        }
        int i = top-1;
        for(;i>=0;i--)
        {
            if(data[i]<=num) break;
        }
        for(int j = top;j>i+1;j--){
            data[j] = data[j-1];
        }
        data[i+1] = num;
        top++;
    }

    public double findMedian() {
        if(top%2==0){
            return (data[top/2]+data[top/2-1])/2.0;
        }
        else{
            return data[top/2];
        }
    }
}

时间复杂度：插入数据时寻找位置O(n) ，移动元素O(n) ，求中位数O(1) 。因此总体的时间复杂度应该是O(n) 。

当然，如此简单的想法，数据量大一点时当然很慢了，我甚至没看到打败了百分之多少。然后我在解答区看到了非常优秀的解法，这里也写出来分享一下。

解法二：两个堆

这种思路是这样的：我维护两片数据，一边存储数字们中较小的一半Tmin，另一边存储数字们中较大的一半Tmax，两边的容量相等或者相差一个，这样中位数就要么是Tmin中最大的数（或Tmax中最小的数，取决于我允许哪一半容量可以大一个），要么是Tmin中的最小与Tmax中的最大的平均值。

于是问题变成了：

  1. 怎么方便的取Tmin中的最大，Tmax中的最小？
  2. 新来一个数据，怎么判断放入Tmin还是Tmax？

对于第一个问题，最适合的数据结构就是大顶堆与小顶堆，维护这两个数据结构可以天然保证Tmin大顶堆的堆顶是其最大值，Tmax小顶堆的堆顶是其最小值。同时这个特性也方便这两片数据平衡容量：

这里把两边分别命名为maxHeap（大顶堆，存小数字）和minHeap（小顶堆，存大数字），并允许maxHeap可以多一个数字。新来一个数字，直接放进maxHeap，然后取出maxHeap中最大的数放入minHeap；若minHeap的容量比maxHeap大，则再取出minHeap中最小的数放入maxHeap。这样就保证了前文所提到的性质，然后获取中位数就很简便了。

由于Java中的堆用于实现了优先队列，而Java中的PriorityQueue默认里是小顶堆，不过也可以在创建PriorityQueue实例时传入一个实现了Comparator接口的类，通过控制比较结果来控制优先级（即可以控制创建大顶堆还是小顶堆）。比如

1
2
3
4
5
6
7
maxHeap = new PriorityQueue<Integer>(new Comparator<Integer>(){
	@Override
	public int compare(Integer o1,Integer o2){
		return o2.compareTo(o1);// return o2 - o1;
	}
});
//maxHeap = new PriorityQueue<Integer>((a,b)->b - a);//简洁写法

优先队列数据结构属于Java集合框架的一部分，关于Java集合框架之后会整理一期博客。

此解法完整代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import java.util.Comparator;
import java.util.PriorityQueue;

class MedianFinder {
    PriorityQueue<Integer> maxHeap; //大顶堆，存小数据
    PriorityQueue<Integer> minHeap; //小顶堆，存大数据

    /** initialize your data structure here. */
    public MedianFinder() {
        this.maxHeap = new PriorityQueue<Integer>(new maxComparator());
        this.minHeap = new PriorityQueue<Integer>(new minComparator());
    }

    public void addNum(int num) {
        maxHeap.add(num);
        minHeap.add(maxHeap.poll());
        if(minHeap.size()>maxHeap.size())
            maxHeap.offer(minHeap.poll());
    }

    public double findMedian() {
        if(minHeap.size() == maxHeap.size())
            return (minHeap.peek()+maxHeap.peek())/2.0;
        else return maxHeap.peek();
    }

    public class minComparator implements Comparator<Integer>
    {
        public int compare(Integer i1, Integer i2)
        {
            return  i1 - i2;
        }
    }

    public class maxComparator implements Comparator<Integer>
    {
        public int compare(Integer i1, Integer i2)
        {
            return  i2 - i1;
        }
    }
}

时间复杂度： 维护堆所需要的时间是O(log n) ，添加数据时只需常数个堆操作。找到平均值也只需常数级复杂度。因此总体的时间复杂度应该是O(log n) 。上面代码在2019/07/11时打败了85.18%的Java，看了看排在前面的代码，思路也是一样的，至于中间用了什么技巧，之后再看吧，目前的主要目的是学习数据结构与算法（可能是少写一个minComparator？）