[最简单实现] 计算 数据流中的中位数 leetcode算法Python实现

一、问题描述：

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

例如，
[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、假如有一个数据流：[3, 1, 9, 25, 12，… ]

每流入一个数据，先在列表尾部append，再做排序，然后计算中位数，过程如下：

3 -> [3]
中位数：3
1 -> [3, 1] -> [1, 3]
中位数：(1+3)/2 = 2.0
9 -> [1, 3, 9] -> [1, 3, 9]
中位数：3
25 -> [1, 3, 9, 25] -> [1, 3, 9, 25]
中位数：(3+9)/2 = 6.0
12 -> [1, 3, 9, 25, 12] -> [1, 3, 9, 12, 25]
中位数：9

最后输出的中位数列表为：[3, 2.0, 3, 6.0, 9, ...]

2、计算过程总结：先排序、再折半

如图所示：

可见，找出中位数分两种情况：

• 当列表长度为奇数：直接取中间数
• 当列表长度为偶数：计算最中间两数的平均数

3、基本实现

算法实现的方法很多，最直观的就是用列表的sort()，再取中值即可。

def running_medians_naive(iterable):
    medians = []
    values = []
    for i,x in enumerate(iterable):
        values.append(x)
        values.sort()
        if i % 2 == 0:
            medians.append(values[i//2])
        else:
            medians.append((values[i//2]+values[i//2+1])/2)
    return medians

running_medians_naive([3, 1, 9, 25, 12])

观察以上算法实现的时间复杂度，可以发现每次都要对列表进行sort()，比较耗时。

4、堆Heap实现

仔细观察计算过程，可以发现：取中位数只跟列表中央的1-2个值相关，如图所示：

跟堆Heap的思想非常相似，可以把中位数的左半折看成是大顶堆max-heap、右半侧看成是小顶堆min-heap，这样就可以用堆Heap来降低时间复杂度了。

Python实现由两部分组成：max-heap堆 和 取中位数算法。

4.1、Max-heap堆

class Heap:
    def __init__(self, key=lambda x:x):
        self.data = []
        self.key  = key

    @staticmethod
    def _parent(idx):
        return (idx-1)//2
        
    @staticmethod
    def _left(idx):
        return idx*2+1

    @staticmethod
    def _right(idx):
        return idx*2+2
    
    def heapify(self, idx=0):
        while idx < len(self.data):
            left = Heap._left(idx)
            right = Heap._right(idx)
            mid_idx = idx
            if left < len(self.data) and self.key(self.data[left]) > self.key(self.data[mid_idx]):
                mid_idx = left
            if right < len(self.data) and self.key(self.data[right]) > self.key(self.data[mid_idx]):
                mid_idx = right
            if mid_idx != idx:
                self.data[mid_idx], self.data[idx] = self.data[idx], self.data[mid_idx]
                idx = mid_idx
            else:
                break
            
    def add(self, x):
        self.data.append(x)
        idx = len(self.data) - 1
        while idx > 0:
            p = Heap._parent(idx)
            if self.key(self.data[idx]) > self.key(self.data[p]):
                self.data[idx], self.data[p] = self.data[p], self.data[idx]
                idx = p
            else:
                break
        
    def peek(self):
        return self.data[0]

    def pop(self):
        ret = self.data[0]
        self.data[0] = self.data[len(self.data)-1]
        del self.data[len(self.data)-1]
        self.heapify()
        return ret
    
    def __bool__(self):
        return len(self.data) > 0

    def __len__(self):
        return len(self.data)

    def __repr__(self):
        return repr(self.data)

4.2、取中位数

def running_medians_heap(iterator):
    medians = []
    max_heap = Heap(lambda x:x)
    min_heap = Heap(lambda x:-x)
    for i, x in enumerate(iterator):
        if i % 2 == 0:
            max_heap.add(x)
            min_heap.add(max_heap.pop())
            medians.append(min_heap.peek())
        else:
            min_heap.add(x)
            max_heap.add(min_heap.pop())
            medians.append((max_heap.peek() + min_heap.peek()) / 2)
    return medians

running_medians_heap([3, 1, 9, 25, 12])

具体算法原理，请参考视频：
https://www.bilibili.com/vide...

转载请注明出处，如有意见或建议，请联系公众号或电邮：bjstorm#foxmail.com