剑指offer-41

剑指 Offer 41. 数据流中的中位数

如何得到一个数据流中的中位数?

如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。

如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

例如,

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构:

  • void addNum(int num) - 从数据流中添加一个整数到数据结构中。
  • double findMedian() - 返回目前所有元素的中位数。

示例:

1
2
3
4
5
6
7
8
9
10
输入:
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出:[null,null,null,1.50000,null,2.00000]

输入:
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出:[null,null,2.00000,null,2.50000]

官方题解

给定一长度为 $N$ 的无序数组,其中位数的计算方法:首先对数组执行排序(使用 $O(N \log N)$ 时间),然后返回中间元素即可(使用 $O(1)$ 时间)。

针对本题,根据以上思路,可以将数据流保存在一个列表中,并在添加元素时 保持数组有序

此方法的时间复杂度为 $O(N)$ ,其中包括:

  • 查找元素插入位置 $O(\log N)$(二分查找)
  • 向数组某位置插入元素 $O(N)$ (插入位置之后的元素都需要向后移动一位)。

借助 可进一步优化时间复杂度。

建立一个 小顶堆 $A$ 和 大顶堆 $B$ ,各保存列表的一半元素,且规定:

  • $A$ 保存 较大 的一半,长度为 $\frac{N}{2}$ ($N$ 为偶数)或 $\frac{N+1}{2}$ ($N$为奇数)
  • $B$ 保存 较小 的一半,长度为 $\frac{N}{2}$ ($N$ 为偶数)或 $\frac{N-1}{2}$ ($N$ 为奇数)

随后,中位数可仅根据 $A, B$ 的堆顶元素计算得到。

image-20210618165657925

算法流程:

设元素总数为 $N = m + n$ ,其中 m 和 n 分别为 A 和 B 中的元素个数。

addNum(num) 函数:

  1. 当 $m==n$ (既$N$为偶数):需要向A添加一个元素。实现方式:将新元素$num$ 加入至 $B$,再将 $B$ 堆定元素插入至 $A$。
  2. 当 $m \neq n$(既 $N$ 为奇数):需要向 $B$ 添加一个元素。实现方式:将新元素$num$加入至 $A$,再将 $A$ 堆顶元素插入至 $B$。

A为小顶堆,存储较大的一半;B为大顶堆,存储较小的一半。

假设插入数字 $num$ 遇到情况 1.。由于 $num$ 可能属于 “较小的一半”(既属于$B$),因此不能将 $nums$ 直接插入至 $A$。而应先将 $num$ 插入至 $B$,再将 $B$ 堆顶元素插入至 $A$。

这样就可以始终保持 $A$ 保存较大的一半、 $B$ 保存较小的一半。

findMedian() 函数:

  1. 当 $m=n$ ($N$ 为偶数):则中位数为($A$的堆顶元素 + $B$的堆顶元素)/2.
  2. 当 $m \neq n$ ($N$为奇数):则中位数为$A$的堆顶元素。

Python 中 heapq 模块是小顶堆。

实现 大顶堆 方法: 小顶堆的插入和弹出操作均将元素 取反 即可。

Java 使用 PriorityQueue<>((x, y) -> (y - x)) 可方便实现大顶堆。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class MedianFinder:

def __init__(self):
"""
initialize your data structure here.
"""
self.A = [] # 小顶堆,保存较大的一半
self.B = [] # 大顶堆,保存较小的一半


def addNum(self, num: int) -> None:
if len(self.A) != len(self.B):
heappush(self.A,num) # 装入元素
heappush(self.B, -heappop(self.A)) # 取出小元素,装入B
else:
heappush(self.B,-num) #
heappush(self.A,-heappop(self.B)) # 取出大元素,装入A


def findMedian(self) -> float:
return self.A[0] if len(self.A) != len(self.B) else (self.A[0] - self.B[0]) / 2.0
作者

bd160jbgm

发布于

2021-06-18

更新于

2021-06-18

许可协议