用大小堆实现插入O(logn),查询O(1)的算法,注意invariant是大堆的size一定是与小堆相等或者比小堆大一。另外大堆存放所有小等于中位数的值,小堆存放所有大于中位数的值。
package Hard;
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
/**
* Numbers are randomly generated and passed to a method. Write a program to find and maintain the median value as new values are generated.
译文:
随机产生一些数传递给一个函数,写程序找出并维护这些数的中位数。
*
*/
public class S18_9 {
static class MaxHeapComparator implements Comparator<Integer>{
// Comparator that sorts integers from highest to lowest
@Override
public int compare(Integer o1, Integer o2) {
return o2-o1;
}
}
static class MinHeapComparator implements Comparator<Integer>{
// Comparator that sorts integers from lowest to highest
@Override
public int compare(Integer o1, Integer o2) {
return o1-o2;
}
}
private static Comparator<Integer> maxHeapComparator;
private static Comparator<Integer> minHeapComparator;
private static PriorityQueue<Integer> maxHeap; // 大堆里放小于中位数的元素
private static PriorityQueue<Integer> minHeap; // 小堆里放大于中位数的元素
// O(logn)
public static void addNewNumber(int randomNumber) {
/* Note: addNewNumber maintains a condition that maxHeap.size() >= minHeap.size() */
if (maxHeap.size() == minHeap.size()) { // 当小堆size与大堆相等时,中位数等于(maxHeap.top()+minHeap.top())/2
if ((minHeap.peek() != null) && randomNumber > minHeap.peek()) { // 新数比中位数大,所以要放入小堆,但因为大堆的size必须大等于小堆
maxHeap.offer(minHeap.poll()); // 所以要先把小堆中最小的元素移到大堆中
minHeap.offer(randomNumber); // 然后再把新数放入小堆
} else { // 新数比中位数小或相等,所以放大堆
maxHeap.offer(randomNumber);
}
}
else { // 这种情况必定是大堆的size比小堆多1,则中位数等于maxHeap.top()
if(randomNumber < maxHeap.peek()){ // 新数比中位数小或相等,所以放大堆
minHeap.offer(maxHeap.poll()); // 先把大堆的最大元素转移到小堆
maxHeap.offer(randomNumber); // 然后把新数放入大堆
}
else { // 新数比中位数大,所以放小堆,因为现在大堆的size比小堆多1,所以可以直接放入小堆,无需转移
minHeap.offer(randomNumber);
}
}
}
// O(1)
public static double getMedian() {
/* maxHeap is always at least as big as minHeap. So if maxHeap is empty, then minHeap is also. */
if (maxHeap.isEmpty()) {
return 0;
}
if (maxHeap.size() == minHeap.size()) { // 当大小堆size相同时,取平均
return ((double)minHeap.peek() + (double) maxHeap.peek()) / 2;
} else { // 否则大堆size比小堆多1,因此中位数就是大堆的最大值
/* If maxHeap and minHeap are of different sizes, then maxHeap must have one extra element. Return maxHeap’s top element.*/
return maxHeap.peek();
}
}
public static void addNewNumberAndPrintMedian(int randomNumber) {
addNewNumber(randomNumber);
System.out.println("Random Number = " + randomNumber);
printMinHeapAndMaxHeap();
System.out.println("\nMedian = " + getMedian() + "\n");
}
public static void printMinHeapAndMaxHeap(){
Integer[] minHeapArray = minHeap.toArray(new Integer[minHeap.size()]);
Integer[] maxHeapArray = maxHeap.toArray(new Integer[maxHeap.size()]);
Arrays.sort(minHeapArray, maxHeapComparator);
Arrays.sort(maxHeapArray, maxHeapComparator);
System.out.print("MinHeap =");
for (int i = minHeapArray.length - 1; i >= 0 ; i--){
System.out.print(" " + minHeapArray[i]);
}
System.out.print("\nMaxHeap =");
for (int i = 0; i < maxHeapArray.length; i++){
System.out.print(" " + maxHeapArray[i]);
}
}
public static void main(String[] args) {
int arraySize = 10;
int range = 7;
maxHeapComparator = new MaxHeapComparator();
minHeapComparator = new MinHeapComparator();
maxHeap = new PriorityQueue<Integer>(arraySize - arraySize/2, maxHeapComparator);
minHeap = new PriorityQueue<Integer>(arraySize/2, minHeapComparator);
for(int i = 0; i < arraySize; i++) {
int randomNumber = (int) (Math.random() * (range+1));
addNewNumberAndPrintMedian(randomNumber);
}
}
}