[程序员代码面试指南]9-输出数据流中位数(堆)

题意

随着数据流的读入，输出中位数，若当前度入数据有偶数个则取中间两个平均值。

题解

大根堆维护小的一半数，小根堆维护大的一半数。
插入时间复杂度O(logn),查找时间复杂度O(1)。

相关

Java中优先队列是用堆实现的。

todo

compare的返回值控制比较、排序还是搞不太懂==。

代码

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Scanner;

class MaxHeapComparator implements Comparator<Integer>{
    @Override
    public int compare(Integer num1,Integer num2) {
        return num2-num1;
    }
}

class MinHeapComparator implements Comparator<Integer>{
    @Override
    public int compare(Integer num1,Integer num2) {
        return num1-num2;
    }
}

class Median{
    PriorityQueue<Integer> maxHeap=new PriorityQueue<>(new MaxHeapComparator());
    PriorityQueue<Integer> minHeap=new PriorityQueue<>(new MinHeapComparator());
    public void add(int num) {
        if(maxHeap.isEmpty()||num<=maxHeap.peek()) {//
            maxHeap.add(num);
        }
        else {
            minHeap.add(num);
        }
        
        if(Math.abs(maxHeap.size()-minHeap.size())>1) {
            if(maxHeap.size()>minHeap.size()) {
                minHeap.add(maxHeap.poll());
            }
            else {
                maxHeap.add(minHeap.poll());
            }
        }
    }
    
    public Integer getMedian() {
        if(maxHeap.isEmpty()) {
            return null;
        }
        if(maxHeap.size()==minHeap.size()) {
            return (maxHeap.peek()+minHeap.peek())/2;
        }
        return maxHeap.size()>minHeap.size()?maxHeap.peek():minHeap.peek();
    }
}

public class Main {
    public static void main(String args[]) {
        Scanner in=new Scanner(System.in);
        Median median=new Median();
        while(in.hasNext()) {
            int num=in.nextInt();
            median.add(num);
            System.out.println(median.getMedian());
        }
    }
}