第63题——数据流中的中位数

450 阅读3分钟

题目:

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。我们使用Insert()方法读取数据流,使用GetMedian()方法获取当前读取数据的中位数。

思路:

先用java集合PriorityQueue来设置一个小顶堆和大顶堆
主要的思想是:因为要求的是中位数,那么这两个堆,大顶堆用来存较小的数,从大到小排列;
小顶堆存较大的数,从小到大的顺序排序,显然中位数就是大顶堆的根节点与小顶堆的根节点和的平均数。
⭐保证:小顶堆中的元素都大于等于大顶堆中的元素,所以每次塞值,并不是直接塞进去,而是从另一个堆中poll出一个最大(最小)的塞值
⭐当数目为偶数的时候,将这个值插入大顶堆中,再将大顶堆中根节点(即最大值)插入到小顶堆中;
⭐当数目为奇数的时候,将这个值插入小顶堆中,再讲小顶堆中根节点(即最小值)插入到大顶堆中;
⭐取中位数的时候,如果当前个数为偶数,显然是取小顶堆和大顶堆根结点的平均值;如果当前个数为奇数,显然是取小顶堆的根节点
理解了上面所述的主体思想,下面举个例子辅助验证一下。

例如,传入的数据为:[5,2,3,4,1,6,7,0,8],那么按照要求,输出是"5.00 3.50 3.00 3.50 3.00 3.50 4.00 3.50 4.00 "

那么整个程序的执行流程应该是(用min表示小顶堆,max表示大顶堆):

5先进入大顶堆,然后将大顶堆中最大值放入小顶堆中,此时min=[5],max=[无],avg=[5.00]
2先进入小顶堆,然后将小顶堆中最小值放入大顶堆中,此时min=[5],max=[2],avg=[(5+2)/2]=[3.50]
3先进入大顶堆,然后将大顶堆中最大值放入小顶堆中,此时min=[3,5],max=[2],avg=[3.00]
4先进入小顶堆,然后将小顶堆中最小值放入大顶堆中,此时min=[4,5],max=[3,2],avg=[(4+3)/2]=[3.50]
1先进入大顶堆,然后将大顶堆中最大值放入小顶堆中,此时min=[3,4,5],max=[2,1],avg=[3/00]
6先进入小顶堆,然后将小顶堆中最小值放入大顶堆中,此时min=[4,5,6],max=[3,2,1],avg=[(4+3)/2]=[3.50]
7先进入大顶堆,然后将大顶堆中最大值放入小顶堆中,此时min=[4,5,6,7],max=[3,2,1],avg=[4]=[4.00]
0先进入小顶堆,然后将小顶堆中最大值放入小顶堆中,此时min=[4,5,6,7],max=[3,2,1,0],avg=[(4+3)/2]=[3.50]
8先进入大顶堆,然后将大顶堆中最小值放入大顶堆中,此时min=[4,5,6,7,8],max=[3,2,1,0],avg=[4.00]

Java

package nowcoder;

import java.util.Comparator;
import java.util.PriorityQueue;

public class S63_GetMedian {
    //小顶堆
    private PriorityQueue<Integer> minHeap = new PriorityQueue<Integer>();
    //大顶堆
    private PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(15, new Comparator<Integer>() {
        @Override
        public int compare(Integer o1, Integer o2) {
            return o2 - o1;
        }
    });
    int count = 0;
    public void insert(Integer num){
        //count为偶数时,先进大根堆,将大根堆的堆顶放到小根堆
        if ((count & 1) == 0) {
            if (!maxHeap.isEmpty() && num < maxHeap.peek()) {
                maxHeap.offer(num);
                num = maxHeap.poll();
            }
            minHeap.offer(num);
        }
        else {
            if (!minHeap.isEmpty() && num > minHeap.peek()){
                minHeap.offer(num);
                num = minHeap.poll();
            }
            maxHeap.offer(num);
        }
        count++;
    }
    public double getMedian(){
        if(count==0)
            throw new RuntimeException("no available number!");
        if ((count & 1) == 0)
            return new Double(minHeap.peek() + maxHeap.peek())/2;
        else
            return new Double(minHeap.peek());
    }
    public static void main(String[] args){
        S63_GetMedian s63 = new S63_GetMedian();
        Integer[] array = {5,2,3,4,1,6,7,0,8};
        for (int i=0;i<array.length;i++){
            s63.insert(array[i]);
            System.out.println(s63.getMedian());
        }
    }
}

Python

import heapq
class GetMedian:
    def __init__(self):
        self.heaps = [], []
    def insert(self, num):
        small, large = self.heaps
        heapq.heappush(small, -heapq.heappushpop(large, num))
        if len(large) < len(small):
            heapq.heappush(large, -heapq.heappop(small))
    def getMedian(self):
        small, large = self.heaps
        if len(large) > len(small):
            return float(large[0])
        return (large[0] - small[0]) / 2.0

if __name__ == '__main__':
    test = GetMedian()
    array = [5, 2, 3, 4, 1, 6, 7, 0, 8]
    for i in array:
        test.insert(i)
        print(test.getMedian())