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

转载

mob604756fadec0 2021-07-13 09:11:00

文章标签 中位数数据结构存储数据持续更新 文章分类 大数据

1. 题目

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

例如，

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

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

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

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

2. 示例

示例1：

输入：
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出：[null,null,null,1.50000,null,2.00000]

示例2：

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

3. 题解

最初的想法是用TreeMap来存储数据，key为输入，value为数量。然后在查找的时候查找偶数查找第n/2个数以及n/2 + 1个数。超时。

然后使用大小堆的方式：

定义一个小顶堆和一个大顶堆。小顶堆存储大的一半，大顶堆存储小的一半。
A保存较大的一半，长度为N/2(N 为偶数) 或 (N + 1) / 2 (N 为奇数);
B保存较小的一半，长度为N/2(N 为偶数) 或 (N - 1) / 2 (N 为奇数)；

4. 实现

4.1 TreeMap（超时）

 1 class MedianFinder {
 2 
 3     Map<Integer, Integer> ans;
 4     int len = 0;
 5 
 6     /** initialize your data structure here. */
 7     public MedianFinder() {
 8         ans =  new TreeMap<>();
 9     }
10 
11     public void addNum(int num) {
12         ans.put(num, ans.getOrDefault(num, 0) + 1);
13         len++;
14     }
15 
16     public double findMedian() {
17         int sum = 0;
18         if(len % 2 == 1) {
19             int index = len / 2 + 1;
20             for(int k : ans.keySet()) {
21                 if(sum + ans.get(k) >= index) {
22                     return k;
23                 }
24                 sum += ans.get(k);
25             }
26         } else {
27             int pre = len / 2, cur = len / 2 + 1;
28             System.out.println("输入：" + pre + " " + cur);
29             int preValue = -50, curValue = -50;
30             for(int k : ans.keySet()) {
31                 if(sum + ans.get(k) >= pre && preValue == -50) {
32                     preValue = k;
33                 }
34                 if(sum + ans.get(k) >= cur && curValue == -50) {
35                     curValue = k;
36                     break;
37                 }
38                 sum += ans.get(k);
39             }
40             System.out.println("输出：" + preValue + " " +curValue);
41             return (double)((preValue + curValue) / 2.0);
42         }
43         return 0;
44     }
45 }
46 
47 /**
48  * Your MedianFinder object will be instantiated and called as such:
49  * MedianFinder obj = new MedianFinder();
50  * obj.addNum(num);
51  * double param_2 = obj.findMedian();
52  */

View Code

4.2 大小顶堆

 1 class MedianFinder {
 2     // 大小堆
 3     Queue<Integer> A, B;
 4     public MedianFinder() {
 5         A = new PriorityQueue<>(); // 小顶堆，保存较大的一堆
 6         B = new PriorityQueue<>((x, y) -> (y - x)); // 大顶堆，保存较小的一半
 7         B = new PriorityQueue<>((x, y) -> (y - x)); // 大顶堆，保存较小的一半
 8     }
 9     // 建立大小顶堆，各保存一半：
10     // A保存较大的一半，长度为N/2(N 为偶数) 或 (N + 1) / 2 (N 为奇数);
11     // B保存较小的一半，长度为N/2(N 为偶数) 或 (N - 1) / 2 (N 为奇数)
12     public void addNum(int num) {
13         // A.size() != B.size()，需向B添加一个元素。实现方法：将新元素num插入A，再将A堆顶元素插入至B；
14         // A.size() == B.size()，需向A添加一个元素，实现方法：将新元素num插入B，再将B堆顶元素插入至A；
15         if(A.size() != B.size()) {
16             A.add(num);
17             B.add(A.poll());
18         } else {
19             B.add(num);
20             A.add(B.poll());
21         }
22     }
23     public double findMedian() {
24         return A.size() != B.size() ? A.peek() : (A.peek() + B.peek()) / 2.0;
25     }
26 }
27 
28 /**
29  * Your MedianFinder object will be instantiated and called as such:
30  * MedianFinder obj = new MedianFinder();
31  * obj.addNum(num);
32  * double param_2 = obj.findMedian();
33  */

View Code