剑指63-数据流中的中位数

Posted on | In |

题目描述

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

  • Insert()方法维护一个排序列表，每读出一个数，用二分查找来找到其位置并插入队列。
    • 每次迭代 mid = (i+j)//2， 根据mid位置的元素与要插入的num的大小关系，更新i或者j。更新确保不破坏循环条件i<=j，还要[i,j]闭区间缩小，以确保不会有死循环。
  • GetMedian()方法在上述排序列表上获取中位数。
  • 运算符优先级： 幂运算，正负号， 算数运算符（— = * /）， 比较运算符（< > <= >=）， 逻辑运算符（not and or）

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
class Solution:
    def __init__(self):
        self.array = []
    def Insert(self, num):
        # write code here
        i , j = 0, len(self.array)-1
        # 使用了二分查找来插入元素
        if i > j:
            self.array.insert(0, num)
            return
        while i <= j:
            if i == j:
                if self.array[i] > num:
                    self.array.insert(i, num)
                else:
                    self.array.insert(i+1, num)
                return
            mid = (i+j)>>1 # 整除
            if self.array[mid] < num: # 三种情况对i和j的处理都不会破坏循环条件 i<=j
                if i == mid: # 确保i和j在每次迭代中范围都会缩小
                    i = mid + 1
                else:
                    i = mid
            elif self.array[mid] > num: # 确保i不会大于j
                j = mid   
            else:  # 这种情况，直接插入肯定是没问题的
                self.array.insert(mid, num)
                return
    # 注意本题的bug：data是一个没有用的参数，但是必须用，不然会出错 
    def GetMedian(self, data):
        # write code here
        mylen = len(self.array)
        if mylen&1 == 1:
            return self.array[mylen>>1]
        else:
            return (self.array[(mylen>>1) - 1] + self.array[mylen//2])/2.0 # 除数必须是小数，不然2.7版本会把它结果转换成整数

考虑如下思路：

  • 维护两个数据结构Less和Largue，每次收到一个数字：
    • 如果该数字是第奇数个，要与Large里的最小数比较，如果较小则直接加入Less；否则取出Large里最小数添加进Less，然后把该数添加进Large。
    • 如果该数字是第偶数个，要与Less里的最大数比较，如果较大则直接加入Large；否则取出Less里的最大数添加进Large，然后把该数添加进Less。
  • 获取中位数时：
    • 如果一共有奇数个数，则返回Less的最大值。
    • 如果一共有偶数个数，则返回Less最大值和Large最小值的平均数。

堆（heap），它是一种优先队列，能够以任意顺序添加对象，并随时找出（并删除）最小的元素（最小堆）。相比于列表方法min，这样做的效率要高得多。Python没有独立的堆类型，而只有一个包含一些堆操作函数的模块，名为heapq。将元素全部取负，用堆来处理，实际上就是用最小堆实现了最大堆。

  • 上面的Less是个最大堆，而Large是个最小堆。
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
import heapq as hq
class Solution:
    def __init__(self):
        self.less = [] # 充当最大堆，存储较小数
        self.large = [] # 最小堆，存储较大数
        self.count = 0
    def Insert(self, num):
        # write code here
        self.count += 1
        if self.count&1 == 1: # 目标是往最大堆里Less添加元素
            min_large_list = hq.nsmallest(1, self.large) 
            if len(min_large_list) == 0 or num <= min_large_list[0]:
                hq.heappush(self.less, -num)
            else: # num比最小堆的最小数还要大
                min_large = hq.heappop(self.large)
                hq.heappush(self.less, -min_large)
                hq.heappush(self.large, num)
        else: # 目标是往最小堆Large里添加元素
            max_less = -hq.nsmallest(1, self.less)[0] # 最大堆的最大值,肯定存在
            if num >= max_less:
                hq.heappush(self.large, num)
            else:
                hq.heappush(self.large, -hq.heappop(self.less))
                hq.heappush(self.less, -num)
        
        
    # 注意本题的bug：data是一个没有用的参数，但是必须用，不然会出错 
    def GetMedian(self):
        # write code here
        if self.count == 0:
            return None
        if self.count&1 == 1:
            return -hq.nsmallest(1, self.less)[0]
        else:
            return (-hq.nsmallest(1, self.less)[0]+hq.nsmallest(1, self.large)[0])/2.0

假如忽略掉题设中数据流这个场景，一开始就能拿到所有数据，则可以用partition函数法，

  • 如果数组长度是奇数，则找到下标为len(nums)>>1的元素将其作为返回值。
  • 如果数组长度是偶数，先找到下标为len(nums)>>1的元素，然后再找到左子序列的最大值，返回二者平均数。