题目
有一个源源不断往外吐出整数的数据流,假设你有足够的空间来保存吐出的数。请设计一个方法,这个方法可以随时取出之前吐出所有数的中位数
要求
如果已经保存了吐出的N个数,那么任意时刻将一个新数加入的过程,其时间复杂度不超过O(logN)取得中位数的过程,时间复杂度为O(1)
思路
建立一个大根堆,一个小根堆。每次来的一个数,和大根堆的堆顶比较,如果小于大根堆的堆顶,就加入大根堆;如果大于大根堆的堆顶,就加入小根堆
同时还要满足这两个堆中的元素个数之差不能超过2(即<2)。例如大根堆中的元素现在有3个,小根堆中的元素有1个,此时就需要把大根堆的堆顶弹出,放入小根堆中;反之也一样。注意:每次往堆中加入数的同时,也要调整堆的结构
如果吐出的数据个数为偶数,则中位数是两个堆的堆顶相加除以2;为奇数,中位数是元素个数较多的那个堆的堆顶
往堆里加入一个数的时间复杂度是O(logN),取出中位数的时间复杂度是O(1),满足题目要求
代码
import java
.util
.Comparator
;
import java
.util
.PriorityQueue
;
public class GetMedian {
public static class MinHeapComparator implements Comparator<Integer> {
@Override
public int compare(Integer a
, Integer b
) {
return a
> b
? 1 : -1;
}
}
public static class MaxHeapComparator implements Comparator<Integer> {
@Override
public int compare(Integer a
, Integer b
) {
return b
> a
? 1 : -1;
}
}
public static void main(String
[] args
) {
int[] arr
= { 8, 3, 4, 6, 9, 7, 1 };
double median
= getMedian(arr
);
System
.out
.println(median
);
}
private static double getMedian(int[] arr
) {
PriorityQueue
<Integer> big
= new PriorityQueue<Integer>(new MaxHeapComparator());
PriorityQueue
<Integer> small
= new PriorityQueue<Integer>(new MinHeapComparator());
for (int i
= 0; i
< arr
.length
; i
++) {
int cur
= arr
[i
];
if (big
.isEmpty() || cur
< big
.peek())
big
.add(cur
);
else
small
.add(cur
);
if (big
.size() - small
.size() >= 2)
small
.add(big
.poll());
else if (small
.size() - big
.size() >= 2)
big
.add(small
.poll());
}
if (arr
.length
% 2 == 0)
return ((double) (big
.peek() + small
.peek()) / 2);
else
return big
.size() > small
.size() ? big
.peek() : small
.peek();
}
}