Sui Xin's Blog

[剑指Offer]数据流中的中位数

Posted on In Disqus:

题目描述

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

解题思路

  1. 方法一：不考虑复杂度。每次取中位数时都对数组进行排序。
  2. 方法二：最大堆/最小堆（剑指Offer答案）。使用最大堆/最小堆的插入时间复杂度为$O(\log n)$，得到中位数的时间复杂度为$O(1)$。原理：
    1. 保证最大堆和最小堆的数据平衡（其中数据的数目之差不大于1）。总数目为偶数时，插入最小堆，为奇数时，插入最大堆；
    2. 保证最大堆中所有元素都小于最小堆中元素。如，总数目为偶数时，应插入最小堆，但该元素小于最大堆堆顶元素，所以反而将该元素插入最大堆，再将最大堆堆顶元素插入最小堆。总数目为奇数时类似；
    3. 取中位数时，若总数目为偶数，则中位数为两个堆堆顶的平均数，为奇数时，中位数为最小堆堆顶元素。

代码

Python(2.7.3)

方法一：不考虑复杂度

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# -*- coding:utf-8 -*-
class Solution:
    def __init__(self):
        self.data = []
    
    def Insert(self, num):
        # write code here
        self.data.append(num)

    def GetMedian(self, _):
        # write code here
        self.data.sort()
        length = len(self.data)
        if length % 2 == 0:
            return (self.data[length / 2 - 1] + self.data[length / 2]) / 2.0
        return self.data[length // 2]

运行时间：21ms
占用内存：5724k

方法二：最大堆/最小堆

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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
# -*- coding:utf-8 -*-
class Solution:
    def __init__(self):
        self.maxHeap = []
        self.minHeap = []
    
    def insertMaxHeap(self, num):
        # 新元素插到堆尾
        self.maxHeap.append(num)
        idx = len(self.maxHeap) - 1
        # 从堆尾开始调整堆，元素不断上浮，直到满足最大堆条件
        while idx > 0:
            if self.maxHeap[(idx - 1) // 2] < self.maxHeap[idx]:
                self.maxHeap[(idx - 1) // 2], self.maxHeap[idx] = self.maxHeap[idx], self.maxHeap[(idx - 1) // 2]
                idx = (idx - 1) // 2
            else:
                break
    
    def popMaxHeap(self):
        # 取出堆顶元素，堆尾元素放到堆顶
        root = self.maxHeap[0]
        self.maxHeap[0] = self.maxHeap[-1]
        self.maxHeap.pop()
        idx = 0
        length = len(self.maxHeap)
        # 从堆顶开始调整堆，元素不断下沉，直到满足最大堆条件
        while 2 * idx + 1 < length:
            change_node = 2 * idx + 1
            if 2 * idx + 2 < length and self.maxHeap[2 * idx + 1] < self.maxHeap[2 * idx + 2]:
                change_node += 1
            if self.maxHeap[idx] < self.maxHeap[change_node]:
                self.maxHeap[idx], self.maxHeap[change_node] = self.maxHeap[change_node], self.maxHeap[idx]
                idx = change_node
            else:
                break
        return root
    
    def insertMinHeap(self, num):
        self.minHeap.append(num)
        idx = len(self.minHeap) - 1
        while idx > 0:
            if self.minHeap[(idx - 1) // 2] > self.minHeap[idx]:
                self.minHeap[(idx - 1) // 2], self.minHeap[idx] = self.minHeap[idx], self.minHeap[(idx - 1) // 2]
                idx = (idx - 1) // 2
            else:
                break
    
    def popMinHeap(self):
        root = self.minHeap[0]
        self.minHeap[0] = self.minHeap[-1]
        self.minHeap.pop()
        idx = 0
        length = len(self.minHeap)
        while 2 * idx + 1 < length:
            change_node = 2 * idx + 1
            if 2 * idx + 2 < length and self.minHeap[2 * idx + 1] > self.minHeap[2 * idx + 2]:
                change_node += 1
            if self.minHeap[idx] > self.minHeap[change_node]:
                self.minHeap[idx], self.minHeap[change_node] = self.minHeap[change_node], self.minHeap[idx]
                idx = change_node
            else:
                break
        return root
    
    def Insert(self, num):
        # write code here
        if (len(self.maxHeap) + len(self.minHeap)) % 2 == 0:
            # 数据总共有偶数个时，新元素插入最小堆
            if len(self.maxHeap) > 0 and num < self.maxHeap[0]:
                # 若新元素小于最大堆堆顶元素，则将该元素插入最大堆，拿出最大堆堆顶元素插入最小堆
                self.insertMaxHeap(num)
                num = self.popMaxHeap()
            self.insertMinHeap(num)
        else:
            # 数据总共有奇数个时，新元素插入最大堆
            if len(self.minHeap) > 0  and num > self.minHeap[0]:
                # 若新元素大于最小堆堆顶元素，则将该元素插入最小堆，拿出最小堆堆顶元素插入最大堆
                self.insertMinHeap(num)
                num = self.popMinHeap()
            self.insertMaxHeap(num)
    
    def GetMedian(self, _):
        # write code here
        if (len(self.maxHeap) + len(self.minHeap)) % 2 == 0:
            return (self.maxHeap[0] + self.minHeap[0]) / 2.0
        else:
            return self.minHeap[0]

运行时间：25ms
占用内存：5724k

Related Posts